feat: Add reconnection capability for persisted WebDriver sessions

src/client.rs

@@ -60,22 +60,40 @@ impl Client {
     /// should generally be one that supports HTTPS, as that is commonly required by WebDriver
     /// implementations.
     ///
-    /// Calls `with_capabilities_and_connector` with an empty capabilities list.
     pub(crate) async fn new_with_connector<C>(
         webdriver: &str,
         connector: C,
     ) -> Result<Self, error::NewSessionError>
     where
         C: connect::Connect + Unpin + 'static + Clone + Send + Sync,
     {
-        Self::with_capabilities_and_connector(
-            webdriver,
-            &webdriver::capabilities::Capabilities::new(),
-            connector,
+        let (client, wdb) = Session::create_client_and_parse_url(webdriver, connector).await?;
+        Session::setup_session(
+            client,
+            wdb,
+            None,
+            Some(webdriver::capabilities::Capabilities::new()),
         )
         .await
     }
 
+    /// Reconnect to a previously established WebDriver session using its ID.
+    ///
+    /// Ideal for resuming operations without losing session data after a disconnect
+    /// or process restart, ensuring that the session can be reused without creating a new one.
+    ///
+    pub async fn with_existing_session<C>(
+        webdriver: &str,
+        session_id: &str,
+        connector: C,
+    ) -> Result<Self, error::NewSessionError>
+    where
+        C: connect::Connect + Unpin + 'static + Clone + Send + Sync,
+    {
+        let (client, wdb) = Session::create_client_and_parse_url(webdriver, connector).await?;
+        Session::setup_session(client, wdb, Some(session_id), None).await
+    }
+
     /// Connect to the WebDriver host running the given address.
     ///
     /// Prefer using [`ClientBuilder`](crate::ClientBuilder) over calling this method directly.

src/session.rs

@@ -535,6 +535,24 @@ impl<C> Session<C>
 where
     C: connect::Connect + Unpin + 'static + Clone + Send + Sync,
 {
+    fn new(
+        rx: mpsc::UnboundedReceiver<Task>,
+        client: hyper_util::client::legacy::Client<C, BoxBody<hyper::body::Bytes, Infallible>>,
+        wdb_url: url::Url,
+        session_id: &str,
+    ) -> Self {
+        Session {
+            ongoing: Ongoing::None,
+            rx,
+            client,
+            wdb: wdb_url,
+            session: Some(session_id.to_string()),
+            is_legacy: false,
+            ua: None,
+            persist: false,
+        }
+    }
+
     fn shutdown(&mut self, ack: Option<Ack>) {
         // session was not created
         if self.session.is_none() {
@@ -607,45 +625,70 @@ where
         }
     }
 
-    pub(crate) async fn with_capabilities_and_connector(
+    pub(crate) async fn create_client_and_parse_url(
         webdriver: &str,
-        cap: &webdriver::capabilities::Capabilities,
         connector: C,
-    ) -> Result<Client, error::NewSessionError> {
+    ) -> Result<
+        (
+            hyper_util::client::legacy::Client<C, BoxBody<hyper::body::Bytes, Infallible>>,
+            url::Url,
+        ),
+        error::NewSessionError,
+    > {
         // Where is the WebDriver server?
-        let wdb = webdriver.parse::<url::Url>();
-        let wdb = wdb.map_err(error::NewSessionError::BadWebdriverUrl)?;
+        let wdb = webdriver
+            .parse::<url::Url>()
+            .map_err(error::NewSessionError::BadWebdriverUrl)?;
+
         // We want a tls-enabled client
         let client = hyper_util::client::legacy::Client::builder(TokioExecutor::new())
             .build::<_, BoxBody<hyper::body::Bytes, Infallible>>(connector);
 
-        let mut cap = cap.to_owned();
+        Ok((client, wdb))
+    }
+
+    pub(crate) async fn setup_session(
+        client: hyper_util::client::legacy::Client<C, BoxBody<hyper::body::Bytes, Infallible>>,
+        wdb: url::Url,
+        session_id: Option<&str>,
+        cap: Option<webdriver::capabilities::Capabilities>,
+    ) -> Result<Client, error::NewSessionError> {
         // We're going to need a channel for sending requests to the WebDriver host
         let (tx, rx) = mpsc::unbounded_channel();
 
-        // Set up our WebDriver session.
-        tokio::spawn(Session {
-            rx,
-            ongoing: Ongoing::None,
-            client,
-            wdb,
-            session: None,
-            is_legacy: false,
-            ua: None,
-            persist: false,
-        });
+        if let Some(session_id) = session_id {
+            // Reconnect to an existing session
+            tokio::spawn(Session::new(rx, client, wdb, session_id));
+        } else if let Some(_capabilities) = cap {
+            // Spawn a new WebDriver session.
+            tokio::spawn(Session {
+                rx,
+                ongoing: Ongoing::None,
+                client,
+                wdb,
+                session: None,
+                is_legacy: false,
+                ua: None,
+                persist: false,
+            });
+        }
 
         // now that the session is running, let's do the handshake
-        let client = Client {
+        Ok(Client {
             tx: tx.clone(),
             is_legacy: false,
             new_session_response: None,
-        };
+        })
+    }
+
+    pub(crate) async fn with_capabilities_and_connector(
+        webdriver: &str,
+        cap: &webdriver::capabilities::Capabilities,
+        connector: C,
+    ) -> Result<Client, error::NewSessionError> {
+        let (client, wdb) = Self::create_client_and_parse_url(webdriver, connector).await?;
+        let mut cap = cap.to_owned();
 
-        // Create a new session for this client
-        // https://www.w3.org/TR/webdriver/#dfn-new-session
-        // https://www.w3.org/TR/webdriver/#capabilities
-        //  - we want the browser to wait for the page to load
         if !cap.contains_key("pageLoadStrategy") {
             cap.insert("pageLoadStrategy".to_string(), Json::from("normal"));
         }
@@ -659,6 +702,12 @@ where
                 .insert("w3c".to_string(), Json::from(true));
         }
 
+        // Create a new session for this client
+        // https://www.w3.org/TR/webdriver/#dfn-new-session
+        // https://www.w3.org/TR/webdriver/#capabilities
+        //  - we want the browser to wait for the page to load
+        let client = Self::setup_session(client, wdb, None, Some(cap.clone())).await?;
+
         let session_config = webdriver::capabilities::SpecNewSessionParameters {
             alwaysMatch: cap.clone(),
             firstMatch: vec![webdriver::capabilities::Capabilities::new()],
@@ -671,7 +720,7 @@ where
             .await
         {
             Ok(new_session_response) => Ok(Client {
-                tx,
+                tx: client.tx,
                 is_legacy: false,
                 new_session_response: Some(wd::NewSessionResponse::from_wd(new_session_response)),
             }),
@@ -718,7 +767,7 @@ where
                     .await?;
 
                 Ok(Client {
-                    tx,
+                    tx: client.tx,
                     is_legacy: true,
                     new_session_response: None,
                 })