From cd6cc5618d48bf8e02df5854a1b276a14bd07b09 Mon Sep 17 00:00:00 2001
From: shrinishLT <shrinishv@lambdatest.com>
Date: Thu, 17 Sep 2026 17:16:15 +0530
Subject: [PATCH] fix: unwrap decorated drivers before reading the session id

Since 1.0.7 (#13) smartuiSnapshot casts the driver to RemoteWebDriver to read the
session id. A driver wrapped with Selenium's EventFiringDecorator is a proxy that
implements JavascriptExecutor and WrapsDriver but is not a RemoteWebDriver, so the
cast throws ClassCastException inside the try block and the only output is
"SmartUI snapshot failed". No snapshot is ever posted to the CLI.

Unwrap through WrapsDriver / Decorated until a RemoteWebDriver appears; send the
snapshot without a session id when none is found. smartuiResults uses the same
helper. The catch block now logs the exception so the next failure explains itself.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
---
 .../io/github/lambdatest/SmartUIResults.java  |  6 ++++-
 .../io/github/lambdatest/SmartUISnapshot.java | 10 +++++----
 .../github/lambdatest/utils/SmartUIUtil.java  | 22 +++++++++++++++++++
 3 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/src/main/java/io/github/lambdatest/SmartUIResults.java b/src/main/java/io/github/lambdatest/SmartUIResults.java
index e837d7c..1141c27 100644
--- a/src/main/java/io/github/lambdatest/SmartUIResults.java
+++ b/src/main/java/io/github/lambdatest/SmartUIResults.java
@@ -44,7 +44,11 @@ public static JSONObject smartuiResults(WebDriver driver) throws Exception {
 
         try {
             // Extract sessionId from the driver (null-safe)
-            org.openqa.selenium.remote.SessionId sid = ((org.openqa.selenium.remote.RemoteWebDriver) driver).getSessionId();
+            org.openqa.selenium.remote.RemoteWebDriver remoteDriver = SmartUIUtil.unwrapRemoteWebDriver(driver);
+            if (remoteDriver == null) {
+                throw new IllegalArgumentException("Driver must be a RemoteWebDriver, or wrap one, to extract sessionId");
+            }
+            org.openqa.selenium.remote.SessionId sid = remoteDriver.getSessionId();
             if (sid == null) {
                 throw new IllegalStateException("Unable to get sessionId from the driver");
             }
diff --git a/src/main/java/io/github/lambdatest/SmartUISnapshot.java b/src/main/java/io/github/lambdatest/SmartUISnapshot.java
index 0edcf5c..8b1daf3 100644
--- a/src/main/java/io/github/lambdatest/SmartUISnapshot.java
+++ b/src/main/java/io/github/lambdatest/SmartUISnapshot.java
@@ -61,9 +61,11 @@ public static JSONObject smartuiSnapshot(WebDriver driver, String snapshotName,
                 ((JavascriptExecutor) driver).executeScript(domString);
 
                 // Append sessionId to options
-                String sessionId = ((org.openqa.selenium.remote.RemoteWebDriver) driver).getSessionId().toString();
-                if (!sessionId.isEmpty()) {
-                    options.put("sessionId", sessionId);
+                org.openqa.selenium.remote.RemoteWebDriver remoteDriver = SmartUIUtil.unwrapRemoteWebDriver(driver);
+                if (remoteDriver != null && remoteDriver.getSessionId() != null) {
+                    options.put("sessionId", remoteDriver.getSessionId().toString());
+                } else {
+                    log.fine("Driver is not a RemoteWebDriver; snapshot sent without sessionId");
                 }
 
                 // Resolve any WebElement objects in element/ignoreDOM/selectDOM to CSS selectors
@@ -143,7 +145,7 @@ public static JSONObject smartuiSnapshot(WebDriver driver, String snapshotName,
             }
 
         } catch (Exception e) {
-            log.severe(String.format(Constants.Errors.SMARTUI_SNAPSHOT_FAILED, snapshotName));
+            log.severe(Constants.Errors.SMARTUI_SNAPSHOT_FAILED + " for '" + snapshotName + "': " + e);
             return null;
         }
     }
diff --git a/src/main/java/io/github/lambdatest/utils/SmartUIUtil.java b/src/main/java/io/github/lambdatest/utils/SmartUIUtil.java
index 62d88ea..d8b6cff 100644
--- a/src/main/java/io/github/lambdatest/utils/SmartUIUtil.java
+++ b/src/main/java/io/github/lambdatest/utils/SmartUIUtil.java
@@ -7,6 +7,10 @@
 import io.github.lambdatest.models.*;
 import com.google.gson.Gson;
 import io.github.lambdatest.constants.Constants;
+import org.openqa.selenium.WebDriver;
+import org.openqa.selenium.WrapsDriver;
+import org.openqa.selenium.remote.RemoteWebDriver;
+import org.openqa.selenium.support.decorators.Decorated;
 
 
 public class SmartUIUtil {
@@ -30,6 +34,24 @@ public SmartUIUtil(String proxyProtocol, String proxyHost, int proxyPort, boolea
         this.httpClient = new HttpClientUtil(proxyProtocol, proxyHost, proxyPort, allowInsecure);
     }
 
+    // Selenium decorators (EventFiringDecorator) hand back a proxy that is not a RemoteWebDriver; unwrap until the real one appears
+    public static RemoteWebDriver unwrapRemoteWebDriver(WebDriver driver) {
+        Object current = driver;
+        for (int depth = 0; current != null && depth < 8; depth++) {
+            if (current instanceof RemoteWebDriver) {
+                return (RemoteWebDriver) current;
+            }
+            if (current instanceof WrapsDriver) {
+                current = ((WrapsDriver) current).getWrappedDriver();
+            } else if (current instanceof Decorated) {
+                current = ((Decorated<?>) current).getOriginal();
+            } else {
+                return null;
+            }
+        }
+        return null;
+    }
+
     public boolean isSmartUIRunning() {
         try {
             httpClient.isSmartUIRunning();