From aae831e1e5516c2fe123586022d60aa6d7fcda9c Mon Sep 17 00:00:00 2001
From: Robert Wojciechowski <hello@robertvoy.com>
Date: Sat, 4 Apr 2026 00:05:49 +0000
Subject: [PATCH] Fix stale master callback port selection

---
 tests/test_network_helpers.py | 18 +++++++++++++++++-
 utils/network.py              | 32 ++++++++++++++++++++++++--------
 2 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/tests/test_network_helpers.py b/tests/test_network_helpers.py
index 749986e..a9d6954 100644
--- a/tests/test_network_helpers.py
+++ b/tests/test_network_helpers.py
@@ -90,8 +90,24 @@ class NetworkHelpersTests(unittest.TestCase):
             "https://master.example.com",
         )
 
+    def test_build_master_url_ignores_stale_saved_port_and_uses_runtime_port(self):
+        cfg = {"master": {"host": "192.168.68.56", "port": 8001}}
+        prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8188)
+        self.assertEqual(
+            network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
+            "http://192.168.68.56:8188",
+        )
+
+    def test_build_master_url_keeps_explicit_port_in_host(self):
+        cfg = {"master": {"host": "192.168.68.56:8001"}}
+        prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8188)
+        self.assertEqual(
+            network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
+            "http://192.168.68.56:8001",
+        )
+
     def test_build_master_url_falls_back_to_server_address(self):
-        cfg = {"master": {"host": ""}}
+        cfg = {"master": {"host": "", "port": 8001}}
         prompt_server = types.SimpleNamespace(address="0.0.0.0", port=8190)
         self.assertEqual(
             network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
diff --git a/utils/network.py b/utils/network.py
index ce5c7a7..4668881 100644
--- a/utils/network.py
+++ b/utils/network.py
@@ -66,6 +66,25 @@ def normalize_host(value):
     return host.split("/")[0]
 
 
+def _split_host_and_port(host):
+    if not host:
+        return host, None
+
+    if host.startswith("["):
+        match = re.match(r"^(\[[^\]]+\])(?::(\d+))?$", host)
+        if match:
+            parsed_port = int(match.group(2)) if match.group(2) else None
+            return match.group(1), parsed_port
+        return host, None
+
+    if host.count(":") == 1:
+        candidate_host, candidate_port = host.rsplit(":", 1)
+        if candidate_port.isdigit():
+            return candidate_host, int(candidate_port)
+
+    return host, None
+
+
 def build_worker_url(worker, endpoint=""):
     """Construct the worker base URL with optional endpoint."""
     host = (worker.get("host") or "").strip()
@@ -126,12 +145,7 @@ def build_master_url(config=None, prompt_server_instance=None):
     prompt_server_instance = prompt_server_instance or server.PromptServer.instance
     master_cfg = (config or {}).get("master", {}) or {}
     configured_host = (master_cfg.get("host") or "").strip()
-    configured_port = master_cfg.get("port")
-    default_port = getattr(prompt_server_instance, "port", 8188) or 8188
-    try:
-        port = int(configured_port or default_port)
-    except (TypeError, ValueError):
-        port = int(default_port)
+    runtime_port = getattr(prompt_server_instance, "port", 8188) or 8188
 
     def _needs_https(hostname):
         hostname = hostname.lower()
@@ -149,10 +163,11 @@ def build_master_url(config=None, prompt_server_instance=None):
         if configured_host.startswith(("http://", "https://")):
             return configured_host.rstrip("/")
 
-        host = configured_host
+        host, explicit_port = _split_host_and_port(configured_host)
+        port = explicit_port if explicit_port is not None else int(runtime_port)
         scheme = "https" if _needs_https(host) or port == 443 else "http"
         default_port_for_scheme = 443 if scheme == "https" else 80
-        if configured_port is None and scheme == "https" and _needs_https(host):
+        if explicit_port is None and scheme == "https" and _needs_https(host):
             port = default_port_for_scheme
         port_part = "" if port == default_port_for_scheme else f":{port}"
         return f"{scheme}://{host}{port_part}"
@@ -160,6 +175,7 @@ def build_master_url(config=None, prompt_server_instance=None):
     address = getattr(prompt_server_instance, "address", "127.0.0.1") or "127.0.0.1"
     if address in ("0.0.0.0", "::"):
         address = "127.0.0.1"
+    port = int(runtime_port)
     scheme = "https" if port == 443 else "http"
     default_port_for_scheme = 443 if scheme == "https" else 80
     port_part = "" if port == default_port_for_scheme else f":{port}"