update

app.py

@@ -38,24 +38,41 @@ def discover_available_models(
     base_url: str = "http://localhost",
     port_range: Tuple[int, int] = (8000, 8010),
 ) -> List[Tuple[str, str]]:
-    """Discover trace eval servers by pinging /health on ports. Returns [(server_url, model_name), ...]."""
+    """Discover trace eval servers by pinging /health. Returns [(server_url, model_name), ...].
+    For ngrok or https URLs, uses the URL as-is. For localhost, scans ports."""
+    base_url = base_url.strip().rstrip("/")
+    urls_to_check: List[Tuple[str, str]] = []
+
+    # Single URL mode: ngrok, https, or URL that already has a port
+    if "ngrok" in base_url or base_url.startswith("https://"):
+        urls_to_check = [(base_url, "Trace (ngrok/external)")]
+    elif ":" in base_url.split("//")[-1].split("/")[0]:
+        # Already has port (e.g. http://localhost:8000)
+        urls_to_check = [(base_url, "Trace")]
+    else:
+        # Scan ports for localhost
+        start_port, end_port = port_range
+        for port in range(start_port, end_port + 1):
+            urls_to_check.append((f"{base_url}:{port}", f"Trace @ port {port}"))
+
     available = []
-    start_port, end_port = port_range
-    for port in range(start_port, end_port + 1):
-        server_url = f"{base_url.rstrip('/')}:{port}"
-        print(f"Checking {server_url}/health")
+    headers = {}
+    if "ngrok" in base_url:
+        headers["ngrok-skip-browser-warning"] = "true"
+    for server_url, label in urls_to_check:
         try:
-            r = requests.get(f"{server_url}/health", timeout=2.0)
+            r = requests.get(f"{server_url}/health", timeout=5.0, headers=headers)
             if r.status_code == 200:
                 try:
-                    info = requests.get(f"{server_url}/model_info", timeout=2.0).json()
-                    print(info)
-                    name = info.get("model_id", f"Trace @ port {port}")
+                    info = requests.get(
+                        f"{server_url}/model_info", timeout=5.0, headers=headers
+                    ).json()
+                    name = info.get("model_id", label)
                 except Exception:
-                    name = f"Trace @ port {port}"
+                    name = label
                 available.append((server_url, name))
-        except requests.exceptions.RequestException:
-            print(f"Error checking {server_url}/health")
+        except requests.exceptions.RequestException as e:
+            logger.debug(f"Could not reach {server_url}/health: {e}")
             continue
     return available
 
@@ -64,8 +81,9 @@ def get_model_info_for_url(server_url: str) -> Optional[str]:
     """Get formatted model info for a trace eval server."""
     if not server_url:
         return None
+    headers = {"ngrok-skip-browser-warning": "true"} if "ngrok" in server_url else {}
     try:
-        r = requests.get(f"{server_url.rstrip('/')}/model_info", timeout=5.0)
+        r = requests.get(f"{server_url.rstrip('/')}/model_info", timeout=5.0, headers=headers)
         if r.status_code == 200:
             return format_trace_model_info(r.json())
     except Exception as e:
@@ -90,8 +108,9 @@ def check_server_health(server_url: str) -> Tuple[str, Optional[dict], Optional[
     """Check trace eval server health. Returns (status_msg, health_data, model_info_text)."""
     if not server_url:
         return "Please provide a server URL.", None, None
+    headers = {"ngrok-skip-browser-warning": "true"} if "ngrok" in server_url else {}
     try:
-        r = requests.get(f"{server_url.rstrip('/')}/health", timeout=5.0)
+        r = requests.get(f"{server_url.rstrip('/')}/health", timeout=5.0, headers=headers)
         r.raise_for_status()
         data = r.json()
         info = get_model_info_for_url(server_url)
@@ -109,10 +128,12 @@ def run_inference_via_server(
     """Run inference via trace eval server. Returns (prediction, overlay_path, trace_points_text)."""
     with open(image_path, "rb") as f:
         image_b64 = base64.b64encode(f.read()).decode("utf-8")
+    headers = {"ngrok-skip-browser-warning": "true"} if "ngrok" in server_url else {}
     r = requests.post(
         f"{server_url.rstrip('/')}/predict",
         json={"image_base64": image_b64, "instruction": instruction},
         timeout=120.0,
+        headers=headers,
     )
     r.raise_for_status()
     data = r.json()