Update bot/integrations/diag_extra.py

bot/integrations/diag_extra.py

@@ -1,14 +1,52 @@
 # PATH: bot/integrations/diag_extra.py
+from __future__ import annotations
+
 import socket
 from urllib.parse import urlparse
 
-def dns_check(url: str) -> str:
+
+def _to_host(url_or_host: str) -> str:
+    s = (url_or_host or "").strip()
+    if not s:
+        return ""
+    # If it's a URL, extract hostname; else treat as hostname
+    if "://" in s:
+        try:
+            u = urlparse(s)
+            return (u.hostname or "").strip()
+        except Exception:
+            return ""
+    # remove port if present
+    if ":" in s:
+        return s.split(":", 1)[0].strip()
+    return s
+
+
+def dns_check(url_or_host: str, max_ips: int = 5) -> str:
+    """
+    Returns a short string:
+      - OK:host -> ip1,ip2
+      - ERR:<Reason>
+    Works in HF container (best-effort).
+    """
+    host = _to_host(url_or_host)
+    if not host:
+        return "ERR:empty_host"
+
     try:
-        host = urlparse(url).hostname
-        if not host:
-            return "bad_url"
+        # getaddrinfo may return duplicates; we de-dupe
         infos = socket.getaddrinfo(host, 443, type=socket.SOCK_STREAM)
-        ips = sorted({i[4][0] for i in infos})
-        return "OK " + ", ".join(ips[:6])
+        ips = []
+        for info in infos:
+            ip = info[4][0]
+            if ip not in ips:
+                ips.append(ip)
+            if len(ips) >= max_ips:
+                break
+        if not ips:
+            return f"ERR:no_ip:{host}"
+        return f"OK:{host} -> {', '.join(ips)}"
+    except socket.gaierror as e:
+        return f"ERR:gaierror:{e}"
     except Exception as e:
         return f"ERR:{type(e).__name__}:{e}"