Browser Stream Refactor Fast Follow (#4070)

skyvern/forge/sdk/routes/streaming/channels/execution.py

@@ -16,6 +16,7 @@ import typing as t
 from contextlib import asynccontextmanager
 
 import structlog
+from playwright.async_api import Browser, BrowserContext, Page, Playwright
 
 from skyvern.forge.sdk.routes.streaming.channels.cdp import CdpChannel
 
@@ -26,6 +27,11 @@ LOG = structlog.get_logger()
 
 
 class ExecutionChannel(CdpChannel):
+    # Explicitly declare inherited attributes for mypy when follow_imports = skip
+    browser: Browser | None
+    browser_context: BrowserContext | None
+    page: Page | None
+    pw: Playwright | None
     """
     ExecutionChannel.
     """
@@ -44,7 +50,7 @@ class ExecutionChannel(CdpChannel):
         }
         """
 
-        selected_text = await self.evaluate_js(js_expression, self.page)
+        selected_text = await self.evaluate_js(js_expression)
 
         if isinstance(selected_text, str) or selected_text is None:
             LOG.info(

skyvern/forge/sdk/routes/streaming/channels/message.py

@@ -121,7 +121,7 @@ class MessageChannel:
             return base | {"browser_session_id": self.browser_session.persistent_browser_session_id}
 
         if self.workflow_run:
-            return base | {"workflow_run_id": self.workflow_run.id}
+            return base | {"workflow_run_id": self.workflow_run.workflow_run_id}
 
         return base
 
@@ -174,8 +174,8 @@ class MessageChannel:
                 datums.append(data)
             except asyncio.TimeoutError:
                 break
-            except RuntimeError:
-                if "not connected" in str(RuntimeError).lower():
+            except RuntimeError as ex:
+                if "not connected" in str(ex).lower():
                     break
             except Exception:
                 LOG.exception(f"{self.class_name} Failed to receive message from message channel", **self.identity)

skyvern/forge/sdk/routes/streaming/messages.py

@@ -82,6 +82,7 @@ async def messages(
             browser_session_id=browser_session_id,
             workflow_run_id=workflow_run_id,
         )
+        await websocket.close(code=1002)
         return
 
     message_channel: MessageChannel