browser sessions v2 - backend (#4515)

Signed-off-by: Benji Visser <benji@093b.org>

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

@@ -281,6 +281,43 @@ async def ask_for_clipboard(vnc_channel: VncChannel) -> None:
         LOG.exception(f"{class_name} Failed to ask for clipboard via CDP", **vnc_channel.identity)
 
 
+# TODO(benji): I hate this function. It's messy and gross. Once we remove v1,
+# we should clean this up.
+def _build_vnc_url_from_browser_address(browser_address: str) -> str | None:
+    """
+    Build a routed VNC URL from a V2 K8s routed browser_address.
+
+    V2 K8s routed browser_address format:
+        wss://{domain}/{session_id}/{token}/devtools/browser/{browser_id}
+
+    Returns VNC URL in format:
+        wss://{domain}/vnc/{session_id}/{token}
+
+    Returns None if browser_address is not a V2 routed URL.
+    """
+    if not browser_address:
+        return None
+
+    parsed = urlparse(browser_address)
+
+    # Check if this looks like a V2 routed URL (wss:// with token in path)
+    if parsed.scheme not in ("wss", "ws"):
+        return None
+
+    # Parse path: /{session_id}/{token}/devtools/browser/{browser_id}
+    path_parts = parsed.path.strip("/").split("/")
+    if len(path_parts) < 4 or path_parts[2] != "devtools":
+        return None
+
+    session_id = path_parts[0]
+    token = path_parts[1]
+    domain = parsed.netloc
+
+    # Build VNC URL with same domain and token
+    scheme = "wss" if parsed.scheme == "wss" else "ws"
+    return f"{scheme}://{domain}/vnc/{session_id}/{token}"
+
+
 async def loop_stream_vnc(vnc_channel: VncChannel) -> None:
     """
     Actually stream the VNC data between a frontend and a browser.
@@ -292,24 +329,28 @@ async def loop_stream_vnc(vnc_channel: VncChannel) -> None:
     browser_session = vnc_channel.browser_session
     class_name = vnc_channel.class_name
 
-    if browser_session:
-        if browser_session.ip_address:
-            if ":" in browser_session.ip_address:
-                ip, _ = browser_session.ip_address.split(":")
-                vnc_url = f"ws://{ip}:{vnc_channel.vnc_port}"
-            else:
-                vnc_url = f"ws://{browser_session.ip_address}:{vnc_channel.vnc_port}"
-        else:
-            browser_address = browser_session.browser_address
-
-            parsed_browser_address = urlparse(browser_address)
-            host = parsed_browser_address.hostname
-            vnc_url = f"ws://{host}:{vnc_channel.vnc_port}"
-    else:
+    if not browser_session:
         raise Exception(f"{class_name} No browser session associated with vnc channel.")
 
-    # NOTE(jdo:streaming-local-dev)
-    # vnc_url = "ws://localhost:6080"
+    # First, check if this is a V2 K8s routed session by examining browser_address
+    # V2 sessions have browser_address like: wss://{domain}/{session_id}/{token}/devtools/...
+    # For these, we need to route VNC through the same nginx proxy
+    routed_vnc_url = _build_vnc_url_from_browser_address(browser_session.browser_address)
+    if routed_vnc_url:
+        vnc_url = routed_vnc_url
+    elif browser_session.ip_address:
+        # V1 ECS sessions: Direct IP connection (ip_address is a public/reachable IP)
+        if ":" in browser_session.ip_address:
+            ip, _ = browser_session.ip_address.split(":")
+            vnc_url = f"ws://{ip}:{vnc_channel.vnc_port}"
+        else:
+            vnc_url = f"ws://{browser_session.ip_address}:{vnc_channel.vnc_port}"
+    else:
+        # Last resort: parse browser_address hostname
+        browser_address = browser_session.browser_address
+        parsed_browser_address = urlparse(browser_address)
+        host = parsed_browser_address.hostname
+        vnc_url = f"ws://{host}:{vnc_channel.vnc_port}"
 
     LOG.info(
         f"{class_name} Connecting to vnc url.",
@@ -317,7 +358,12 @@ async def loop_stream_vnc(vnc_channel: VncChannel) -> None:
         **vnc_channel.identity,
     )
 
-    async with websockets.connect(vnc_url) as novnc_ws:
+    # For routed VNC URLs (wss://), we need to pass the x-api-key header for authentication
+    extra_headers: dict[str, str] = {}
+    if vnc_url.startswith("wss://") and vnc_channel.x_api_key:
+        extra_headers["x-api-key"] = vnc_channel.x_api_key
+
+    async with websockets.connect(vnc_url, additional_headers=extra_headers) as novnc_ws:
 
         async def frontend_to_browser() -> None:
             nonlocal class_name

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

@@ -5,8 +5,10 @@ Provides WS endpoints for streaming messages to/from our frontend application.
 import structlog
 from fastapi import WebSocket
 
+from skyvern.config import settings
 from skyvern.forge.sdk.routes.routers import base_router, legacy_base_router
-from skyvern.forge.sdk.routes.streaming.auth import auth
+from skyvern.forge.sdk.routes.streaming.auth import _auth as local_auth
+from skyvern.forge.sdk.routes.streaming.auth import auth as real_auth
 from skyvern.forge.sdk.routes.streaming.channels.message import (
     Loops,
     MessageChannel,
@@ -60,6 +62,7 @@ async def messages(
     client_id: str | None = None,
     token: str | None = None,
 ) -> None:
+    auth = local_auth if settings.ENV == "local" else real_auth
     organization_id = await auth(apikey=apikey, token=token, websocket=websocket)
 
     if not organization_id:
@@ -98,7 +101,7 @@ async def messages(
         )
     else:
         LOG.error(
-            "Message channel: no browser_session_id or workflow_run_id provided.",
+            "[WS] messages: no browser_session_id or workflow_run_id provided",
             client_id=client_id,
             organization_id=organization_id,
         )
@@ -139,5 +142,4 @@ async def messages(
             workflow_run_id=workflow_run_id,
             organization_id=organization_id,
         )
-
         await message_channel.close(reason="message-stream-closed")

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

@@ -45,12 +45,12 @@ async def verify_browser_session(
     """
     Verify the browser session exists, and is usable.
     """
-
     if settings.ENV == "local":
         dummy_browser_session = AddressablePersistentBrowserSession(
             persistent_browser_session_id=browser_session_id,
             organization_id=organization_id,
             browser_address="0.0.0.0:9223",
+            ip_address="localhost",
             created_at=datetime.now(),
             modified_at=datetime.now(),
         )
@@ -199,6 +199,7 @@ async def verify_workflow_run(
             persistent_browser_session_id=workflow_run_id,
             organization_id=organization_id,
             browser_address="0.0.0.0:9223",
+            ip_address="localhost",
             created_at=datetime.now(),
             modified_at=datetime.now(),
         )

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

@@ -12,8 +12,10 @@ NOTE(jdo:streaming-local-dev)
 import structlog
 from fastapi import WebSocket
 
+from skyvern.config import settings
 from skyvern.forge.sdk.routes.routers import base_router, legacy_base_router
-from skyvern.forge.sdk.routes.streaming.auth import auth
+from skyvern.forge.sdk.routes.streaming.auth import _auth as local_auth
+from skyvern.forge.sdk.routes.streaming.auth import auth as real_auth
 from skyvern.forge.sdk.routes.streaming.channels.vnc import (
     Loops,
     VncChannel,
@@ -86,6 +88,7 @@ async def stream(
         workflow_run_id=workflow_run_id,
     )
 
+    auth = local_auth if settings.ENV == "local" else real_auth
     organization_id = await auth(apikey=apikey, token=token, websocket=websocket)
 
     if not organization_id: