[rb] remove bidi callbacks from sockets when unsubscribing

rb/lib/selenium/devtools/support/cdp/domain.rb.erb

@@ -36,7 +36,7 @@ module Selenium
 
         def on(event, &block)
           event = EVENTS[event] if event.is_a?(Symbol)
-          @devtools.callbacks["<%= domain[:domain] %>.#{event}"] << block
+          @devtools.callbacks["<%= domain[:domain] %>.#{event}"][SecureRandom.uuid] = block
         end
 
         <% domain[:commands].each do |command| %>

rb/lib/selenium/webdriver/bidi.rb

@@ -43,8 +43,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/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/network.rbs

@@ -10,7 +10,7 @@ module Selenium
 
         def initialize: (BiDi bidi) -> void
 
-        def add_intercept: (?phases: Array[String], ?contexts: BrowsingContext?, ?url_patterns: String | Array[String]?) -> Hash[String, String]
+        def add_intercept: (?phases: Array[String], ?contexts: BrowsingContext?, ?url_patterns: String | Array[String]?, ?pattern_type: Symbol) -> String
 
         def cancel_auth: -> Hash[nil, nil]
 
@@ -30,7 +30,9 @@ module Selenium
 
         def set_cache_behavior: (String behavior, Array[BrowsingContext]) -> Hash[nil, nil]
 
-        def on: (Symbol event) { (?) -> untyped } -> Hash[nil, nil]
+        def on: (Symbol event) { (?) -> untyped } -> String
+
+        def unsubscribe: ((String | Symbol) event, String id) -> untyped
       end
     end
   end

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

@@ -10,7 +10,7 @@ module Selenium
 
         def status: () -> untyped
 
-        def subscribe: (untyped events, ?untyped? browsing_contexts) -> Hash[nil, nil]
+        def subscribe: (untyped events, ?untyped? browsing_contexts) -> String
 
         def unsubscribe: (untyped events, ?untyped? browsing_contexts) -> untyped
       end

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

@@ -1,31 +1,34 @@
 module Selenium
   module WebDriver
     class Network
+
+      Registration: untyped
+
       @network: BiDi::Network
 
-      @callbacks: Hash[String, String]
+      @registrations: Array[Struct[untyped]]
 
-      attr_reader callbacks: Hash[String, String]
+      attr_reader registrations: Array[Struct[untyped]]
 
       attr_reader network: BiDi::Network
 
       alias bidi network
 
-      def initialize: (Remote::Bridge bridge) -> void
+      def initialize: (Remote::BiDiBridge bridge) -> void
 
-      def remove_handler: (Numeric id) -> untyped
+      def remove_handler: (Struct[untyped] registration) -> untyped
 
       def clear_handlers: () -> untyped
 
-      def add_authentication_handler: (?String? username, ?String? password, *String filter, ?pattern_type: Symbol?) { (?) -> untyped } -> untyped
+      def add_authentication_handler: (?String? username, ?String? password, *String filter, ?pattern_type: Symbol?) ?{ (?) -> untyped } -> Struct[untyped]
 
-      def add_request_handler: (*String filter, ?pattern_type: Symbol?) -> Hash[String, String]
+      def add_request_handler: (*String filter, ?pattern_type: Symbol?) -> Struct[untyped]
 
-      def add_response_handler: (*String filter, ?pattern_type: Symbol?) -> Hash[String, String]
+      def add_response_handler: (*String filter, ?pattern_type: Symbol?) -> Struct[untyped]
 
       private
 
-      def add_handler: (Symbol event_type, String phase, BiDi::InterceptedRequest | BiDi::InterceptedAuth | BiDi::InterceptedResponse intercept_type, Array[String] filter, ?pattern_type: Symbol?) { (untyped) -> untyped } -> untyped
+      def add_handler: [T < BiDi::InterceptedItem] (Symbol event_type, String phase, T intercept_type, Array[String] filter, ?pattern_type: Symbol?) { (untyped) -> untyped } -> Struct[untyped]
     end
   end
 end

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

@@ -35,13 +35,15 @@ module Selenium
 
       def initialize: (url: untyped) -> void
 
-      def add_callback: (untyped event) { () -> void } -> untyped
+      def add_callback: (String event, String id) { () -> void } -> untyped
+
+      def clear_callbacks: -> untyped
 
       def close: () -> untyped
 
       def callbacks: () -> untyped
 
-      def remove_callback: (untyped event, untyped id) -> untyped
+      def remove_callback: (String event, String id) -> untyped
 
       def send_cmd: (**untyped payload) -> untyped
 

rb/spec/integration/selenium/webdriver/bidi/network_spec.rb

@@ -51,7 +51,7 @@ class BiDi
         it 'removes an intercept' do
           network = described_class.new(driver.bidi)
           intercept = network.add_intercept(phases: [described_class::PHASES[:before_request]])
-          expect(network.remove_intercept(intercept['intercept'])).to be_empty
+          expect(network.remove_intercept(intercept)).to be_empty
         end
 
         it 'continues with auth' do