From 531c2d6c99b28136a6a1d2be6a31e2cf90137575 Mon Sep 17 00:00:00 2001
From: batu_hoang <longvatrong111@gmail.com>
Date: Fri, 23 May 2025 16:37:18 +0800
Subject: [PATCH] Fix webdriver wait for response from constellation

Signed-off-by: batu_hoang <longvatrong111@gmail.com>
---
 components/shared/embedder/input_events.rs |  6 ++--
 components/webdriver_server/actions.rs     | 33 ++++++++++++++++++----
 2 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/components/shared/embedder/input_events.rs b/components/shared/embedder/input_events.rs
index 869c4eee004..ff48bd7b7d8 100644
--- a/components/shared/embedder/input_events.rs
+++ b/components/shared/embedder/input_events.rs
@@ -58,16 +58,16 @@ impl InputEvent {
         }
     }
 
-    pub fn with_webdriver_message_id(self, webdriver_id: Option<WebDriverMessageId>) -> Self {
+    pub fn with_webdriver_message_id(mut self, webdriver_id: Option<WebDriverMessageId>) -> Self {
         match self {
             InputEvent::EditingAction(..) => {},
             InputEvent::Gamepad(..) => {},
             InputEvent::Ime(..) => {},
             InputEvent::Keyboard(..) => {},
-            InputEvent::MouseButton(mut event) => {
+            InputEvent::MouseButton(ref mut event) => {
                 event.webdriver_id = webdriver_id;
             },
-            InputEvent::MouseMove(mut event) => {
+            InputEvent::MouseMove(ref mut event) => {
                 event.webdriver_id = webdriver_id;
             },
             InputEvent::Touch(..) => {},
diff --git a/components/webdriver_server/actions.rs b/components/webdriver_server/actions.rs
index b86ce5e8792..0afd113956f 100644
--- a/components/webdriver_server/actions.rs
+++ b/components/webdriver_server/actions.rs
@@ -144,10 +144,33 @@ impl Handler {
             // Step 1.4. Wait for
             // The user agent event loop has spun enough times to process the DOM events
             // generated by the last invocation of the dispatch tick actions steps.
-            //
-            // To ensure we wait for all events to be processed, only the last event in
-            // this tick action step holds the message id.
-            // Whenever a new event is generated, the message id is passed to it.
+            self.wait_for_user_agent_handling_complete(tick_actions)?;
+        }
+
+        // Step 2. Return success with data null.
+        dbg!("Dispatch actions completed successfully");
+        Ok(())
+    }
+
+    fn wait_for_user_agent_handling_complete(
+        &self,
+        tick_actions: &TickActions,
+    ) -> Result<(), ErrorStatus> {
+        // TODO: Add matches! for wheel and key actions
+        // after implmenting webdriver id for wheel and key events.
+        let count_non_null_actions_in_tick = tick_actions
+            .iter()
+            .filter(|(_, action)| {
+                !matches!(action, ActionItem::Pointer(PointerActionItem::Pointer(_)))
+            })
+            .count();
+
+        // To ensure we wait for all events to be processed, only the last event
+        // in each tick action step holds the message id.
+        // Whenever a new event is generated, the message id is passed to it.
+        //
+        // Wait for count_non_null_actions_in_tick number of responses
+        for _ in 0..count_non_null_actions_in_tick {
             match self.constellation_receiver.recv() {
                 Ok(response) => {
                     let current_waiting_id = self
@@ -167,8 +190,6 @@ impl Handler {
             };
         }
 
-        // Step 2. Return success with data null.
-        dbg!("Dispatch actions completed successfully");
         Ok(())
     }