From 4cbaa8979fa58ed74c58e033cb25b480a3a1b198 Mon Sep 17 00:00:00 2001
From: titusfortner <titus.fortner@gmail.com>
Date: Mon, 3 Jun 2024 16:38:41 -0500
Subject: [PATCH] [rb] manage bidi instance on the bridge not the driver

---
 rb/lib/selenium/webdriver/common/driver.rb            |  7 +------
 .../webdriver/common/driver_extensions/has_bidi.rb    |  2 +-
 rb/lib/selenium/webdriver/remote/bridge.rb            | 11 ++++++++++-
 rb/sig/lib/selenium/webdriver/common/driver.rbs       |  1 -
 rb/sig/lib/selenium/webdriver/remote/bridge.rbs       |  1 +
 5 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/rb/lib/selenium/webdriver/common/driver.rb b/rb/lib/selenium/webdriver/common/driver.rb
index 6a816cb67b301..2b1b3c754450c 100644
--- a/rb/lib/selenium/webdriver/common/driver.rb
+++ b/rb/lib/selenium/webdriver/common/driver.rb
@@ -70,7 +70,6 @@ def for(browser, opts = {})
 
       def initialize(bridge: nil, listener: nil, **opts)
         @devtools = nil
-        @bidi = nil
         bridge ||= create_bridge(**opts)
         add_extensions(bridge.browser)
         @bridge = listener ? Support::EventFiringBridge.new(bridge, listener) : bridge
@@ -174,7 +173,6 @@ def quit
       ensure
         @service_manager&.stop
         @devtools&.close
-        @bidi&.close
       end
 
       #
@@ -182,10 +180,7 @@ def quit
       #
 
       def close
-        # If no top-level browsing contexts are open after calling close,
-        # it indicates that the WebDriver session is closed.
-        # If the WebDriver session is closed, the BiDi session also needs to be closed.
-        bridge.close.tap { |handles| @bidi&.close if handles&.empty? }
+        bridge.close
       end
 
       #
diff --git a/rb/lib/selenium/webdriver/common/driver_extensions/has_bidi.rb b/rb/lib/selenium/webdriver/common/driver_extensions/has_bidi.rb
index b0e435d4b6665..859e44737e705 100644
--- a/rb/lib/selenium/webdriver/common/driver_extensions/has_bidi.rb
+++ b/rb/lib/selenium/webdriver/common/driver_extensions/has_bidi.rb
@@ -28,7 +28,7 @@ module HasBiDi
         #
 
         def bidi
-          @bidi ||= Selenium::WebDriver::BiDi.new(url: capabilities[:web_socket_url])
+          @bridge.bidi
         end
       end # HasBiDi
     end # DriverExtensions
diff --git a/rb/lib/selenium/webdriver/remote/bridge.rb b/rb/lib/selenium/webdriver/remote/bridge.rb
index 773d3ef37d778..1d725f4b6325b 100644
--- a/rb/lib/selenium/webdriver/remote/bridge.rb
+++ b/rb/lib/selenium/webdriver/remote/bridge.rb
@@ -213,10 +213,12 @@ def quit
           http.close
         rescue *QUIT_ERRORS
           nil
+        ensure
+          @bidi.close
         end
 
         def close
-          execute :close_window
+          execute(:close_window).tap { |handles| @bidi.close if handles.empty? }
         end
 
         def refresh
@@ -602,6 +604,13 @@ def user_verified(verified, authenticator_id)
           execute :set_user_verified, {authenticatorId: authenticator_id}, {isUserVerified: verified}
         end
 
+        def bidi
+          msg = 'this operation requires enabling BiDi by setting #web_socket_url to true in options class'
+          raise(WebDriver::Error::WebDriverError, msg) unless capabilities.web_socket_url
+
+          @bidi ||= Selenium::WebDriver::BiDi.new(url: capabilities[:web_socket_url])
+        end
+
         def command_list
           COMMANDS
         end
diff --git a/rb/sig/lib/selenium/webdriver/common/driver.rbs b/rb/sig/lib/selenium/webdriver/common/driver.rbs
index 67f2a9017680f..102ad2c526dd2 100644
--- a/rb/sig/lib/selenium/webdriver/common/driver.rbs
+++ b/rb/sig/lib/selenium/webdriver/common/driver.rbs
@@ -5,7 +5,6 @@ module Selenium
 
       include TakesScreenshot
 
-      @bidi: untyped
       @devtools: untyped
       @navigate: untyped
 
diff --git a/rb/sig/lib/selenium/webdriver/remote/bridge.rbs b/rb/sig/lib/selenium/webdriver/remote/bridge.rbs
index a303ec5316006..32aaa1a21b8ef 100644
--- a/rb/sig/lib/selenium/webdriver/remote/bridge.rbs
+++ b/rb/sig/lib/selenium/webdriver/remote/bridge.rbs
@@ -5,6 +5,7 @@ module Selenium
         include _CommandList
         include _Features
 
+        @bidi: WebDriver::BiDi
         @http: untyped
 
         @file_detector: untyped