fix(session): eliminate memory leak in session managers

packages/webdriverio/src/session/context.ts

@@ -24,6 +24,9 @@ export class ContextManager extends SessionManager {
     #mobileContext?: string
     #isNativeContext: boolean
     #getContextSupport = true
+    #onCommandListener: (event: { command: string, body: unknown }) => void
+    #onCommandResultMobileListener: (event: { command: string, result: unknown }) => void
+    #navigationStartedListener: (nav: local.BrowsingContextNavigationInfo) => void
 
     constructor(browser: WebdriverIO.Browser) {
         super(browser, ContextManager.name)
@@ -36,6 +39,10 @@ export class ContextManager extends SessionManager {
             isNativeContext: this.#isNativeContext
         })
 
+        this.#onCommandListener = this.#onCommand.bind(this)
+        this.#onCommandResultMobileListener = this.#onCommandResultMobile.bind(this)
+        this.#navigationStartedListener = this.#navigationStarted.bind(this)
+
         /**
          * Listens for the 'closeWindow' browser command to handle context changes.
          * (classic + bidi)
@@ -53,13 +60,13 @@ export class ContextManager extends SessionManager {
          * Listens for the 'switchToWindow' browser command to handle context changes.
          * Updates the browsingContext with the context passed in 'switchToWindow'.
          */
-        this.#browser.on('command', this.#onCommand.bind(this))
+        this.#browser.on('command', this.#onCommandListener)
 
         /**
          * Listens for the 'closeWindow' browser command to handle context changes.
          */
         if (this.#browser.isMobile) {
-            this.#browser.on('result', this.#onCommandResultMobile.bind(this))
+            this.#browser.on('result', this.#onCommandResultMobileListener)
         } else {
             /**
              * Listen to the 'browsingContext.navigationStarted' event to handle context changes
@@ -68,43 +75,47 @@ export class ContextManager extends SessionManager {
             this.#browser.sessionSubscribe({
                 events: ['browsingContext.navigationStarted']
             })
-            this.#browser.on('browsingContext.navigationStarted', async (nav) => {
-                /**
-                 * no need to do anything as we navigate within the same context
-                 */
-                if (!this.#currentContext || nav.context === this.#currentContext) {
-                    return
-                }
-
-                /**
-                 * a navigation event may have changed the tree structure, so we need to get the
-                 * current tree and see if our context is still there, if not, we need to reset
-                 * the context to the first context in the tree.
-                 */
-                const { contexts } = await this.#browser.browsingContextGetTree({})
-                /**
-                 * check if the context is still in the tree, if not, switch to...
-                 */
-                const hasContext = this.findContext(this.#currentContext, contexts, 'byContextId')
-                /**
-                 * ...the context we are navigating to
-                 */
-                const newContext = contexts.find((context) => context.context === nav.context)
-                if (!hasContext && newContext) {
-                    this.setCurrentContext(newContext.context)
-                    this.#browser.switchToWindow(this.#currentContext)
-                    return
-                }
-            })
+            this.#browser.on('browsingContext.navigationStarted', this.#navigationStartedListener)
         }
     }
 
     removeListeners(): void {
         super.removeListeners()
         // this.#browser.off('result', this.#onCommandResultBidiAndClassic.bind(this))
-        this.#browser.off('command', this.#onCommand.bind(this))
+        this.#browser.off('command', this.#onCommandListener)
         if (this.#browser.isMobile) {
-            this.#browser.off('result', this.#onCommandResultMobile.bind(this))
+            this.#browser.off('result', this.#onCommandResultMobileListener)
+        } else {
+            this.#browser.off('browsingContext.navigationStarted', this.#navigationStartedListener)
+        }
+    }
+
+    async #navigationStarted(nav: local.BrowsingContextNavigationInfo) {
+        /**
+         * no need to do anything as we navigate within the same context
+         */
+        if (!this.#currentContext || nav.context === this.#currentContext) {
+            return
+        }
+
+        /**
+         * a navigation event may have changed the tree structure, so we need to get the
+         * current tree and see if our context is still there, if not, we need to reset
+         * the context to the first context in the tree.
+         */
+        const { contexts } = await this.#browser.browsingContextGetTree({})
+        /**
+         * check if the context is still in the tree, if not, switch to...
+         */
+        const hasContext = this.findContext(this.#currentContext, contexts, 'byContextId')
+        /**
+         * ...the context we are navigating to
+         */
+        const newContext = contexts.find((context) => context.context === nav.context)
+        if (!hasContext && newContext) {
+            this.setCurrentContext(newContext.context)
+            await this.#browser.switchToWindow(this.#currentContext)
+            return
         }
     }
 

packages/webdriverio/src/session/dialog.ts

@@ -15,6 +15,8 @@ export class DialogManager extends SessionManager {
     #initialize: Promise<boolean>
     #autoHandleDialog = true
 
+    #handleUserPromptListener = this.#handleUserPrompt.bind(this)
+
     constructor(browser: WebdriverIO.Browser) {
         super(browser, DialogManager.name)
         this.#browser = browser
@@ -37,12 +39,12 @@ export class DialogManager extends SessionManager {
         this.#browser.on('_dialogListenerRegistered', () => this.#switchListenerFlag(false))
         // @ts-ignore this is a private event
         this.#browser.on('_dialogListenerRemoved', () => this.#switchListenerFlag(true))
-        this.#browser.on('browsingContext.userPromptOpened', this.#handleUserPrompt.bind(this))
+        this.#browser.on('browsingContext.userPromptOpened', this.#handleUserPromptListener)
     }
 
     removeListeners(): void {
         super.removeListeners()
-        this.#browser.off('browsingContext.userPromptOpened', this.#handleUserPrompt.bind(this))
+        this.#browser.off('browsingContext.userPromptOpened', this.#handleUserPromptListener)
         this.#browser.removeAllListeners('_dialogListenerRegistered')
         this.#browser.removeAllListeners('_dialogListenerRemoved')
     }

packages/webdriverio/src/session/networkManager.ts

@@ -20,6 +20,11 @@ export class NetworkManager extends SessionManager {
     #requests = new Map<Context, WebdriverIO.Request>()
     #lastNetworkId?: string
 
+    #navigationStartedListener = this.#navigationStarted.bind(this)
+    #responseCompletedListener = this.#responseCompleted.bind(this)
+    #beforeRequestSentListener = this.#beforeRequestSent.bind(this)
+    #fetchErrorListener = this.#fetchError.bind(this)
+
     constructor(browser: WebdriverIO.Browser) {
         super(browser, NetworkManager.name)
         this.#browser = browser
@@ -44,20 +49,20 @@ export class NetworkManager extends SessionManager {
                 'network.fetchError'
             ]
         }).then(() => true, () => false)
-        this.#browser.on('browsingContext.navigationStarted', this.#navigationStarted.bind(this))
-        this.#browser.on('browsingContext.fragmentNavigated', this.#navigationStarted.bind(this))
-        this.#browser.on('network.responseCompleted', this.#responseCompleted.bind(this))
-        this.#browser.on('network.beforeRequestSent', this.#beforeRequestSent.bind(this))
-        this.#browser.on('network.fetchError', this.#fetchError.bind(this))
+        this.#browser.on('browsingContext.navigationStarted', this.#navigationStartedListener)
+        this.#browser.on('browsingContext.fragmentNavigated', this.#navigationStartedListener)
+        this.#browser.on('network.responseCompleted', this.#responseCompletedListener)
+        this.#browser.on('network.beforeRequestSent', this.#beforeRequestSentListener)
+        this.#browser.on('network.fetchError', this.#fetchErrorListener)
     }
 
     removeListeners(): void {
         super.removeListeners()
-        this.#browser.off('browsingContext.navigationStarted', this.#navigationStarted.bind(this))
-        this.#browser.off('browsingContext.fragmentNavigated', this.#navigationStarted.bind(this))
-        this.#browser.off('network.responseCompleted', this.#responseCompleted.bind(this))
-        this.#browser.off('network.beforeRequestSent', this.#beforeRequestSent.bind(this))
-        this.#browser.off('network.fetchError', this.#fetchError.bind(this))
+        this.#browser.off('browsingContext.navigationStarted', this.#navigationStartedListener)
+        this.#browser.off('browsingContext.fragmentNavigated', this.#navigationStartedListener)
+        this.#browser.off('network.responseCompleted', this.#responseCompletedListener)
+        this.#browser.off('network.beforeRequestSent', this.#beforeRequestSentListener)
+        this.#browser.off('network.fetchError', this.#fetchErrorListener)
     }
 
     async initialize () {

packages/webdriverio/src/session/polyfill.ts

@@ -39,6 +39,8 @@ export class PolyfillManager extends SessionManager {
     #browser: WebdriverIO.Browser
     #scriptsRegisteredInContexts: Set<string> = new Set()
 
+    #registerScriptsListener = this.#registerScripts.bind(this)
+
     constructor(browser: WebdriverIO.Browser) {
         super(browser, PolyfillManager.name)
         this.#browser = browser
@@ -51,6 +53,9 @@ export class PolyfillManager extends SessionManager {
             return
         }
 
+        // start listening for browsingContext.contextCreated
+        this.#browser.on('browsingContext.contextCreated', this.#registerScriptsListener)
+
         /**
          * apply polyfill script for upcoming as well as current execution context
          */
@@ -62,13 +67,12 @@ export class PolyfillManager extends SessionManager {
                 events: ['browsingContext.contextCreated']
             })
         ]).then(() => true, () => false)
-
-        this.#browser.on('browsingContext.contextCreated', this.#registerScripts.bind(this))
     }
 
     removeListeners() {
         super.removeListeners()
-        this.#browser.off('browsingContext.contextCreated', this.#registerScripts.bind(this))
+        // stop listening for browsingContext.contextCreated
+        this.#browser.off('browsingContext.contextCreated', this.#registerScriptsListener)
     }
 
     #registerScripts (context: Pick<local.BrowsingContextInfo, 'context' | 'parent'>) {

packages/webdriverio/src/session/session.ts

@@ -18,11 +18,13 @@ export class SessionManager {
         this.#scope = scope
         const registrationId = `${this.#browser.sessionId}-${this.#scope}`
         if (!listenerRegisteredSession.has(registrationId)) {
-            this.#browser.on('command', this.#onCommand.bind(this))
+            this.#browser.on('command', this.#onCommandListener)
             listenerRegisteredSession.add(registrationId)
         }
     }
 
+    #onCommandListener = this.#onCommand.bind(this)
+
     #onCommand(ev: { command: string }) {
         if (ev.command === 'deleteSession') {
             const sessionManagerInstances = sessionManager.get(this.#scope)
@@ -35,7 +37,7 @@ export class SessionManager {
     }
 
     removeListeners() {
-        this.#browser.off('command', this.#onCommand.bind(this))
+        this.#browser.off('command', this.#onCommandListener)
     }
 
     initialize(): unknown {

packages/webdriverio/src/session/shadowRoot.ts

@@ -24,6 +24,10 @@ export class ShadowRootManager extends SessionManager {
     #documentElement?: remote.ScriptNodeRemoteValue
     #frameDepth = 0
 
+    #handleLogEntryListener = this.handleLogEntry.bind(this)
+    #commandResultHandlerListener = this.#commandResultHandler.bind(this)
+    #handleBidiCommandListener = this.#handleBidiCommand.bind(this)
+
     constructor(browser: WebdriverIO.Browser) {
         super(browser, ShadowRootManager.name)
         this.#browser = browser
@@ -42,19 +46,19 @@ export class ShadowRootManager extends SessionManager {
         this.#initialize = this.#browser.sessionSubscribe({
             events: ['log.entryAdded', 'browsingContext.navigationStarted']
         }).then(() => true, () => false)
-        this.#browser.on('log.entryAdded', this.handleLogEntry.bind(this))
-        this.#browser.on('result', this.#commandResultHandler.bind(this))
-        this.#browser.on('bidiCommand', this.#handleBidiCommand.bind(this))
+        this.#browser.on('log.entryAdded', this.#handleLogEntryListener)
+        this.#browser.on('result', this.#commandResultHandlerListener)
+        this.#browser.on('bidiCommand', this.#handleBidiCommandListener)
         this.#browser.scriptAddPreloadScript({
             functionDeclaration: customElementWrapper.toString()
         })
     }
 
     removeListeners(): void {
         super.removeListeners()
-        this.#browser.off('log.entryAdded', this.handleLogEntry.bind(this))
-        this.#browser.off('result', this.#commandResultHandler.bind(this))
-        this.#browser.off('bidiCommand', this.#handleBidiCommand.bind(this))
+        this.#browser.off('log.entryAdded', this.#handleLogEntryListener)
+        this.#browser.off('result', this.#commandResultHandlerListener)
+        this.#browser.off('bidiCommand', this.#handleBidiCommandListener)
     }
 
     async initialize () {

packages/webdriverio/tests/session/session.test.ts

@@ -1,6 +1,7 @@
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
 import { SessionManager } from '../../src/session/session.js'
+import { ContextManager } from '../../src/session/context.js'
 
 describe('SessionManager', () => {
     const browser ={
@@ -23,4 +24,50 @@ describe('SessionManager', () => {
         new SessionManager(browser, 'dummy')
         expect(browser.on).toHaveBeenCalledTimes(1)
     })
-})
+
+    it('should allow to remove event listeners', () => {
+        const browser = {
+            on: vi.fn(),
+            off: vi.fn(),
+            sessionId: '1234'
+        } as any as WebdriverIO.Browser
+        const sm = new SessionManager(browser, 'scope')
+        const listener = vi.mocked(browser.on).mock.calls[0][1]
+        sm.removeListeners()
+        expect(browser.off).toBeCalledWith('command', listener)
+    })
+
+    it('should remove ContextManager listeners using the same references as they were registered', () => {
+        const browser = {
+            sessionId: '1234',
+            capabilities: {},
+            isBidi: false,
+            isMobile: true,
+            isAndroid: false,
+            on: vi.fn(),
+            off: vi.fn(),
+            sessionSubscribe: vi.fn(),
+            browsingContextGetTree: vi.fn(),
+            switchToWindow: vi.fn(),
+        } as any as WebdriverIO.Browser
+
+        const cm = new ContextManager(browser)
+
+        const onCalls = vi.mocked(browser.on).mock.calls
+        const commandListeners = onCalls.filter(([event]) => event === 'command').map(([, listener]) => listener)
+        const resultListeners = onCalls.filter(([event]) => event === 'result').map(([, listener]) => listener)
+        const baseCommandListener = commandListeners[0]
+        const contextCommandListener = commandListeners[1]
+
+        expect(baseCommandListener).toBeTypeOf('function')
+        expect(contextCommandListener).toBeTypeOf('function')
+        expect(resultListeners).toHaveLength(1)
+        expect(resultListeners[0]).toBeTypeOf('function')
+
+        cm.removeListeners()
+
+        expect(browser.off).toHaveBeenCalledWith('command', baseCommandListener)
+        expect(browser.off).toHaveBeenCalledWith('command', contextCommandListener)
+        expect(browser.off).toHaveBeenCalledWith('result', resultListeners[0])
+    })
+})