[rb] remove bidi callbacks from sockets when unsubscribing

Author: titusfortner

rb/lib/selenium/webdriver/bidi.rb

@@ -44,8 +44,8 @@ def callbacks
         @ws.callbacks
       end
 
-      def add_callback(event, &block)
-        @ws.add_callback(event, &block)
+      def add_callback(event, id, &block)
+        @ws.add_callback(event, id, &block)
       end
 
       def remove_callback(event, id)

rb/lib/selenium/webdriver/bidi/log_handler.rb

@@ -26,19 +26,19 @@ class LogHandler
 
         def initialize(bidi)
           @bidi = bidi
-          @log_entry_subscribed = false
         end
 
         # @return [int] id of the handler
         # steep:ignore:start
         def add_message_handler(type)
-          subscribe_log_entry unless @log_entry_subscribed
-          @bidi.add_callback('log.entryAdded') do |params|
+          id = subscribe_log_entry
+          @bidi.add_callback('log.entryAdded', id) do |params|
             if params['type'] == type
               log_entry_klass = type == 'console' ? ConsoleLogEntry : JavaScriptLogEntry
               yield(log_entry_klass.new(**params))
             end
           end
+          id
         end
         # steep:ignore:end
 
@@ -52,12 +52,10 @@ def remove_message_handler(id)
 
         def subscribe_log_entry
           @bidi.session.subscribe('log.entryAdded')
-          @log_entry_subscribed = true
         end
 
         def unsubscribe_log_entry
-          @bidi.session.unsubscribe('log.entryAdded')
-          @log_entry_subscribed = false
+          @bidi.session.unsubscribe(events: 'log.entryAdded')
         end
       end # LogHandler
     end # Bidi

rb/lib/selenium/webdriver/bidi/log_inspector.rb

@@ -51,7 +51,7 @@ def initialize(driver, browsing_context_ids = nil)
           end
 
           @bidi = driver.bidi
-          @bidi.session.subscribe('log.entryAdded', browsing_context_ids)
+          @browsing_context_ids = browsing_context_ids
         end
 
         def on_console_entry(filter_by = nil, &block)
@@ -96,7 +96,8 @@ def on_log(filter_by = nil, &block)
 
         def on(event, &block)
           event = EVENTS[event] if event.is_a?(Symbol)
-          @bidi.add_callback("log.#{event}", &block)
+          id = @bidi.session.subscribe('log.entryAdded', @browsing_context_ids)
+          @bidi.add_callback("log.#{event}", id, &block)
         end
 
         def check_valid_filter(filter_by)

rb/lib/selenium/webdriver/bidi/network.rb

@@ -45,13 +45,19 @@ def add_intercept(phases: [], contexts: nil, url_patterns: nil, pattern_type: :s
           @bidi.send_cmd('network.addIntercept',
                          phases: phases,
                          contexts: contexts,
-                         urlPatterns: url_patterns)
+                         urlPatterns: url_patterns)['intercept']
         end
 
         def remove_intercept(intercept)
           @bidi.send_cmd('network.removeIntercept', intercept: intercept)
         end
 
+        def unsubscribe(event, id)
+          event = EVENTS[event] if event.is_a?(Symbol)
+          @bidi.session.unsubscribe(ids: id)
+          @bidi.remove_callback(event.to_s, id)
+        end
+
         def continue_with_auth(request_id, username, password)
           @bidi.send_cmd(
             'network.continueWithAuth',
@@ -133,8 +139,9 @@ def set_cache_behavior(behavior, *contexts)
 
         def on(event, &block)
           event = EVENTS[event] if event.is_a?(Symbol)
-          @bidi.add_callback(event, &block)
-          @bidi.session.subscribe(event)
+          id = @bidi.session.subscribe(event)
+          @bidi.add_callback(event, id, &block)
+          id
         end
       end # Network
     end # BiDi

rb/lib/selenium/webdriver/bidi/session.rb

@@ -36,13 +36,13 @@ def subscribe(events, browsing_contexts = nil)
           opts = {events: Array(events)}
           opts[:browsing_contexts] = Array(browsing_contexts) if browsing_contexts
 
-          @bidi.send_cmd('session.subscribe', **opts)
+          @bidi.send_cmd('session.subscribe', **opts)['subscription']
         end
 
-        def unsubscribe(events, browsing_contexts = nil)
-          opts = {events: Array(events)}
-          opts[:browsing_contexts] = Array(browsing_contexts) if browsing_contexts
+        def unsubscribe(events: nil, ids: nil)
+          raise ArgumentError('unsubscribe by event or by id, not both') if events && ids
 
+          opts = events ? {events: Array(events)} : {subscriptions: Array(ids)}
           @bidi.send_cmd('session.unsubscribe', **opts)
         end
       end # Session

rb/lib/selenium/webdriver/common/network.rb

@@ -24,24 +24,26 @@ module WebDriver
     class Network
       extend Forwardable
 
-      attr_reader :callbacks, :network
+      Registration = Struct.new(:subscription, :interception, :event)
+
+      attr_reader :registrations, :network
       alias bidi network
 
       def_delegators :network, :continue_with_auth, :continue_with_request, :continue_with_response
 
       def initialize(bridge)
         @network = BiDi::Network.new(bridge.bidi)
-        @callbacks = {}
+        @registrations = []
       end
 
-      def remove_handler(id)
-        intercept = callbacks[id]
-        network.remove_intercept(intercept['intercept'])
-        callbacks.delete(id)
+      def remove_handler(registration)
+        network.remove_intercept(registration.interception)
+        network.unsubscribe(registration.event, registration.subscription)
+        registrations.delete(registration)
       end
 
       def clear_handlers
-        callbacks.each_key { |id| remove_handler(id) }
+        registrations.dup.each { |registration| remove_handler(registration) }
       end
 
       def add_authentication_handler(username = nil, password = nil, *filter, pattern_type: nil, &block)
@@ -87,15 +89,17 @@ def add_response_handler(*filter, pattern_type: nil, &block)
       private
 
       def add_handler(event_type, phase, intercept_type, filter, pattern_type: nil)
-        intercept = network.add_intercept(phases: [phase], url_patterns: [filter].flatten, pattern_type: pattern_type)
-        callback_id = network.on(event_type) do |event|
+        interception = network.add_intercept(phases: [phase], url_patterns: [filter].flatten,
+                                             pattern_type: pattern_type)
+        subscription = network.on(event_type) do |event|
           request = event['request']
           intercepted_item = intercept_type.new(network, request)
           yield(intercepted_item)
         end
 
-        callbacks[callback_id] = intercept
-        callback_id
+        registration = Registration.new(event: event_type, subscription: subscription, interception: interception)
+        @registrations << registration
+        registration
       end
     end # Network
   end # WebDriver

rb/lib/selenium/webdriver/common/websocket_connection.rb

@@ -49,19 +49,23 @@ def close
       end
 
       def callbacks
-        @callbacks ||= Hash.new { |callbacks, event| callbacks[event] = [] }
+        @callbacks ||= Hash.new { |callbacks, event| callbacks[event] = {} }
       end
 
-      def add_callback(event, &block)
-        callbacks[event] << block
-        block.object_id
+      def add_callback(event, id, &block)
+        callbacks[event][id] = block
       end
 
       def remove_callback(event, id)
-        return if callbacks[event].reject! { |callback| callback.object_id == id }
+        if callbacks.dig(event, id).nil?
+          raise Error::WebDriverError, "Callback with ID #{id} does not exist for event #{event}"
+        end
+
+        callbacks[event].delete(id)
+      end
 
-        ids = callbacks[event]&.map(&:object_id)
-        raise Error::WebDriverError, "Callback with ID #{id} does not exist for event #{event}: #{ids}"
+      def clear_callbacks
+        @callbacks = Hash.new { |callbacks, event| callbacks[event] = {} }
       end
 
       def send_cmd(**payload)
@@ -78,7 +82,7 @@ def send_cmd(**payload)
       private
 
       # We should be thread-safe to use the hash without synchronization
-      # because its keys are WebSocket message identifiers and they should be
+      # because its keys are WebSocket message identifiers, and they should be
       # unique within a devtools session.
       def messages
         @messages ||= {}
@@ -102,7 +106,7 @@ def attach_socket_listener
               next unless message['method']
 
               params = message['params']
-              callbacks[message['method']].each do |callback|
+              callbacks[message['method']].each_value do |callback|
                 @callback_threads.add(callback_thread(params, &callback))
               end
             end

rb/sig/lib/selenium/webdriver/bidi.rbs

@@ -7,13 +7,13 @@ module Selenium
 
       def initialize: (url: String) -> void
 
-      def add_callback: (String | Symbol event) { () -> void } -> Integer
+      def add_callback: (String event, String id) { () -> void } -> Integer
 
       def close: () -> nil
 
       def callbacks: () -> Hash[untyped, untyped]
 
-      def remove_callback: (String event, Integer id) -> Error::WebDriverError?
+      def remove_callback: (String event, String id) -> Error::WebDriverError?
 
       def session: () -> Session
 

rb/sig/lib/selenium/webdriver/bidi/log_handler.rbs

@@ -18,7 +18,7 @@ module Selenium
 
         private
 
-        def subscribe_log_entry: () -> bool
+        def subscribe_log_entry: () -> string
 
         def unsubscribe_log_entry: () -> bool
       end

rb/sig/lib/selenium/webdriver/bidi/log_inspector.rbs

@@ -8,6 +8,8 @@ module Selenium
 
         LOG_LEVEL: Hash[Symbol, String]
 
+        @browsing_context_ids: Array[String]?
+
         def initialize: (untyped driver, ?untyped? browsing_context_ids) -> void
 
         def on_console_entry: (?untyped? filter_by) { () -> untyped } -> untyped