From c7bd726aceaacd1e9d8bd05dfa84a162b921698c Mon Sep 17 00:00:00 2001
From: LawyZheng <lawyzheng1106@gmail.com>
Date: Thu, 17 Jul 2025 17:04:25 +0800
Subject: [PATCH] =?UTF-8?q?dns=20for=20browser=20session=EF=BC=88part=202)?=
 =?UTF-8?q?=20(#2973)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 skyvern/forge/sdk/routes/streaming_vnc.py        | 16 ++++++++++++----
 .../sdk/schemas/persistent_browser_sessions.py   |  1 +
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/skyvern/forge/sdk/routes/streaming_vnc.py b/skyvern/forge/sdk/routes/streaming_vnc.py
index 01b69f79..a9de15f5 100644
--- a/skyvern/forge/sdk/routes/streaming_vnc.py
+++ b/skyvern/forge/sdk/routes/streaming_vnc.py
@@ -169,11 +169,19 @@ async def loop_stream_vnc(streaming: sc.Streaming) -> None:
         LOG.info("No browser session found for task.", task=streaming.task, organization_id=streaming.organization_id)
         return
 
-    browser_address = streaming.browser_session.browser_address
+    vnc_url: str = ""
+    if streaming.browser_session.ip_address:
+        if ":" in streaming.browser_session.ip_address:
+            ip, _ = streaming.browser_session.ip_address.split(":")
+            vnc_url = f"ws://{ip}:{streaming.vnc_port}"
+        else:
+            vnc_url = f"ws://{streaming.browser_session.ip_address}:{streaming.vnc_port}"
+    else:
+        browser_address = streaming.browser_session.browser_address
 
-    parsed_browser_address = urlparse(browser_address)
-    host = parsed_browser_address.hostname
-    vnc_url = f"ws://{host}:{streaming.vnc_port}"
+        parsed_browser_address = urlparse(browser_address)
+        host = parsed_browser_address.hostname
+        vnc_url = f"ws://{host}:{streaming.vnc_port}"
 
     LOG.info(
         "Connecting to VNC URL.",
diff --git a/skyvern/forge/sdk/schemas/persistent_browser_sessions.py b/skyvern/forge/sdk/schemas/persistent_browser_sessions.py
index dba192d0..f0d91ca4 100644
--- a/skyvern/forge/sdk/schemas/persistent_browser_sessions.py
+++ b/skyvern/forge/sdk/schemas/persistent_browser_sessions.py
@@ -11,6 +11,7 @@ class PersistentBrowserSession(BaseModel):
     runnable_type: str | None = None
     runnable_id: str | None = None
     browser_address: str | None = None
+    ip_address: str | None = None
     status: str | None = None
     timeout_minutes: int | None = None
     started_at: datetime | None = None