From c8d249742327d6836ba719a329c94f02c5671313 Mon Sep 17 00:00:00 2001
From: decolua <decoluadt@example.com>
Date: Fri, 20 Mar 2026 12:44:20 +0700
Subject: [PATCH] refactor: update MitmServerCard and MitmToolCard to use
 modalError instead of message for error handling

- Replaced message state with modalError in both components for better error management.
- Removed unused message display logic and adjusted action handling to improve clarity.
- Enhanced error handling in doAction and doDnsAction functions to ignore errors gracefully.
- Updated API call responses to streamline user feedback on actions.
---
 .../cli-tools/components/MitmServerCard.js    | 58 ++++------------
 .../cli-tools/components/MitmToolCard.js      | 68 +++++++------------
 .../api/cli-tools/antigravity-mitm/route.js   |  2 +-
 src/app/api/tunnel/enable/route.js            |  4 ++
 src/lib/tunnel/cloudflared.js                 | 15 +++-
 src/lib/tunnel/tunnelManager.js               | 46 ++++++++++---
 src/mitm/manager.js                           | 11 +++
 src/shared/services/initializeApp.js          | 16 ++++-
 8 files changed, 117 insertions(+), 103 deletions(-)

diff --git a/src/app/(dashboard)/dashboard/cli-tools/components/MitmServerCard.js b/src/app/(dashboard)/dashboard/cli-tools/components/MitmServerCard.js
index 6ab089d3..466fe0c8 100644
--- a/src/app/(dashboard)/dashboard/cli-tools/components/MitmServerCard.js
+++ b/src/app/(dashboard)/dashboard/cli-tools/components/MitmServerCard.js
@@ -13,11 +13,11 @@ export default function MitmServerCard({ apiKeys, cloudEnabled, onStatusChange }
   const [showPasswordModal, setShowPasswordModal] = useState(false);
   const [sudoPassword, setSudoPassword] = useState("");
   const [selectedApiKey, setSelectedApiKey] = useState("");
-  const [message, setMessage] = useState(null);
-  const [pendingAction, setPendingAction] = useState(null); // "start" | "stop"
+  const [pendingAction, setPendingAction] = useState(null);
+  const [modalError, setModalError] = useState(null);
 
   const isWindows = typeof navigator !== "undefined" && navigator.userAgent?.includes("Windows");
-  const isAdmin = status?.isAdmin !== false; // default true until status loaded
+  const isAdmin = status?.isAdmin !== false;
 
   useEffect(() => {
     if (apiKeys?.length > 0 && !selectedApiKey) {
@@ -48,61 +48,39 @@ export default function MitmServerCard({ apiKeys, cloudEnabled, onStatusChange }
     } else {
       setPendingAction(action);
       setShowPasswordModal(true);
-      setMessage(null);
+      setModalError(null);
     }
   };
 
   const doAction = async (action, password) => {
     setLoading(true);
-    setMessage(null);
     try {
       if (action === "trust-cert") {
-        const res = await fetch("/api/cli-tools/antigravity-mitm", {
+        await fetch("/api/cli-tools/antigravity-mitm", {
           method: "PATCH",
           headers: { "Content-Type": "application/json" },
           body: JSON.stringify({ action: "trust-cert", sudoPassword: password }),
         });
-        const data = await res.json();
-        if (res.ok) {
-          setMessage({ type: "success", text: "Certificate trusted successfully" });
-        } else {
-          setMessage({ type: "error", text: data.error || "Failed to trust certificate" });
-        }
       } else if (action === "start") {
         const keyToUse = selectedApiKey?.trim()
           || (apiKeys?.length > 0 ? apiKeys[0].key : null)
           || (!cloudEnabled ? "sk_9router" : null);
-
-        const res = await fetch("/api/cli-tools/antigravity-mitm", {
+        await fetch("/api/cli-tools/antigravity-mitm", {
           method: "POST",
           headers: { "Content-Type": "application/json" },
           body: JSON.stringify({ apiKey: keyToUse, sudoPassword: password }),
         });
-        const data = await res.json();
-        if (res.ok) {
-          setMessage({ type: "success", text: "Server started" });
-        } else {
-          setMessage({ type: "error", text: data.error || "Failed to start server" });
-        }
       } else {
-        const res = await fetch("/api/cli-tools/antigravity-mitm", {
+        await fetch("/api/cli-tools/antigravity-mitm", {
           method: "DELETE",
           headers: { "Content-Type": "application/json" },
           body: JSON.stringify({ sudoPassword: password }),
         });
-        const data = await res.json();
-        if (res.ok) {
-          setMessage({ type: "success", text: "Server stopped — all DNS cleared" });
-        } else {
-          setMessage({ type: "error", text: data.error || "Failed to stop server" });
-        }
       }
       setShowPasswordModal(false);
       setSudoPassword("");
       await fetchStatus();
-    } catch (error) {
-      setMessage({ type: "error", text: error.message });
-    } finally {
+    } catch { /* ignore */ } finally {
       setLoading(false);
       setPendingAction(null);
     }
@@ -110,7 +88,7 @@ export default function MitmServerCard({ apiKeys, cloudEnabled, onStatusChange }
 
   const handleConfirmPassword = () => {
     if (!sudoPassword.trim()) {
-      setMessage({ type: "error", text: "Sudo password is required" });
+      setModalError("Sudo password is required");
       return;
     }
     doAction(pendingAction, sudoPassword);
@@ -159,7 +137,7 @@ export default function MitmServerCard({ apiKeys, cloudEnabled, onStatusChange }
             </p>
           </div>
 
-          {/* API Key selector (only when stopped, to pick key for start) */}
+          {/* API Key selector (only when stopped) */}
           {!isRunning && (
             <div className="flex items-center gap-2">
               <span className="text-xs text-text-muted shrink-0">API Key</span>
@@ -179,16 +157,8 @@ export default function MitmServerCard({ apiKeys, cloudEnabled, onStatusChange }
             </div>
           )}
 
-          {message && (
-            <div className={`flex items-center gap-2 px-2 py-1.5 rounded text-xs ${message.type === "success" ? "bg-green-500/10 text-green-600" : "bg-red-500/10 text-red-600"}`}>
-              <span className="material-symbols-outlined text-[14px]">{message.type === "success" ? "check_circle" : "error"}</span>
-              <span>{message.text}</span>
-            </div>
-          )}
-
-          {/* Action button */}
+          {/* Action buttons */}
           <div className="flex items-center gap-2 flex-wrap" data-i18n-skip="true">
-            {/* Trust Cert button — only when cert exists but not trusted */}
             {status?.certExists && !status?.certTrusted && (
               <button
                 onClick={() => handleAction("trust-cert")}
@@ -249,14 +219,14 @@ export default function MitmServerCard({ apiKeys, cloudEnabled, onStatusChange }
               onChange={(e) => setSudoPassword(e.target.value)}
               onKeyDown={(e) => { if (e.key === "Enter" && !loading) handleConfirmPassword(); }}
             />
-            {message && (
+            {modalError && (
               <div className="flex items-center gap-2 px-2 py-1.5 rounded text-xs bg-red-500/10 text-red-600">
                 <span className="material-symbols-outlined text-[14px]">error</span>
-                <span>{message.text}</span>
+                <span>{modalError}</span>
               </div>
             )}
             <div className="flex items-center justify-end gap-2">
-              <Button variant="ghost" size="sm" onClick={() => { setShowPasswordModal(false); setSudoPassword(""); setMessage(null); }} disabled={loading}>
+              <Button variant="ghost" size="sm" onClick={() => { setShowPasswordModal(false); setSudoPassword(""); setModalError(null); }} disabled={loading}>
                 Cancel
               </Button>
               <Button variant="primary" size="sm" onClick={handleConfirmPassword} loading={loading}>
diff --git a/src/app/(dashboard)/dashboard/cli-tools/components/MitmToolCard.js b/src/app/(dashboard)/dashboard/cli-tools/components/MitmToolCard.js
index 9a9c7dbd..f6b18f46 100644
--- a/src/app/(dashboard)/dashboard/cli-tools/components/MitmToolCard.js
+++ b/src/app/(dashboard)/dashboard/cli-tools/components/MitmToolCard.js
@@ -7,8 +7,6 @@ import Image from "next/image";
 /**
  * Per-tool MITM card — shows DNS status + model mappings.
  * - Auto-saves model mapping on blur or modal select
- * - Start/Stop DNS replaces Save Mappings button
- * - Toggle switch removed; status badge is display-only
  * - Skips sudo modal if password is already cached
  * - Model mappings can only be edited when DNS is active
  */
@@ -27,10 +25,11 @@ export default function MitmToolCard({
   onDnsChange,
 }) {
   const [loading, setLoading] = useState(false);
-  const [message, setMessage] = useState(null);
+  const [warning, setWarning] = useState(null);
   const [showPasswordModal, setShowPasswordModal] = useState(false);
   const [sudoPassword, setSudoPassword] = useState("");
   const [pendingDnsAction, setPendingDnsAction] = useState(null);
+  const [modalError, setModalError] = useState(null);
   const [modelMappings, setModelMappings] = useState({});
   const [modalOpen, setModalOpen] = useState(false);
   const [currentEditingAlias, setCurrentEditingAlias] = useState(null);
@@ -81,7 +80,6 @@ export default function MitmToolCard({
     saveMappings(updated);
   };
 
-  // DNS toggle logic
   const handleDnsToggle = () => {
     if (!serverRunning) return;
     const action = dnsActive ? "disable" : "enable";
@@ -90,13 +88,13 @@ export default function MitmToolCard({
     } else {
       setPendingDnsAction(action);
       setShowPasswordModal(true);
-      setMessage(null);
+      setModalError(null);
     }
   };
 
   const doDnsAction = async (action, password) => {
     setLoading(true);
-    setMessage(null);
+    setWarning(null);
     try {
       const res = await fetch("/api/cli-tools/antigravity-mitm", {
         method: "PATCH",
@@ -107,24 +105,13 @@ export default function MitmToolCard({
       if (!res.ok) throw new Error(data.error || "Failed to toggle DNS");
 
       if (action === "enable") {
-        setMessage({
-          type: "success",
-          text: "DNS enabled successfully.",
-          warning: `Please restart ${tool.name} to apply changes.`,
-        });
-      } else {
-        setMessage({
-          type: "success",
-          text: "DNS disabled — traffic restored",
-        });
+        setWarning(`Restart ${tool.name} to apply changes`);
       }
 
       setShowPasswordModal(false);
       setSudoPassword("");
       onDnsChange?.(data);
-    } catch (error) {
-      setMessage({ type: "error", text: error.message });
-    } finally {
+    } catch { /* ignore */ } finally {
       setLoading(false);
       setPendingDnsAction(null);
     }
@@ -132,7 +119,7 @@ export default function MitmToolCard({
 
   const handleConfirmPassword = () => {
     if (!sudoPassword.trim()) {
-      setMessage({ type: "error", text: "Sudo password is required" });
+      setModalError("Sudo password is required");
       return;
     }
     doDnsAction(pendingDnsAction, sudoPassword);
@@ -185,21 +172,6 @@ export default function MitmToolCard({
               )}
             </div>
 
-            {message && (
-              <div className="flex flex-col gap-1">
-                <div className={`flex items-center gap-2 px-2 py-1.5 rounded text-xs ${message.type === "success" ? "bg-green-500/10 text-green-600" : "bg-red-500/10 text-red-600"}`}>
-                  <span className="material-symbols-outlined text-[14px]">{message.type === "success" ? "check_circle" : "error"}</span>
-                  <span>{message.text}</span>
-                </div>
-                {message.warning && (
-                  <div className="flex items-center gap-2 px-2 py-1.5 rounded text-xs bg-amber-500/10 text-amber-600 border border-amber-500/20">
-                    <span className="material-symbols-outlined text-[14px]">warning</span>
-                    <span className="font-medium">{message.warning}</span>
-                  </div>
-                )}
-              </div>
-            )}
-
             {/* Model Mappings */}
             {tool.defaultModels?.length > 0 && (
               <div className="flex flex-col gap-2">
@@ -245,7 +217,7 @@ export default function MitmToolCard({
             )}
 
             {/* Start / Stop DNS button */}
-            <div>
+            <div className="flex flex-col gap-2 items-start">
               {dnsActive ? (
                 <button
                   onClick={handleDnsToggle}
@@ -256,16 +228,22 @@ export default function MitmToolCard({
                   Stop DNS
                 </button>
               ) : (
-                <Button
-                  variant="primary"
-                  size="sm"
+                <button
                   onClick={handleDnsToggle}
-                  loading={loading}
                   disabled={!serverRunning || loading}
+                  className="px-4 py-1.5 rounded-lg bg-primary/10 border border-primary/30 text-primary font-medium text-xs flex items-center gap-1.5 hover:bg-primary/20 transition-colors disabled:opacity-50 disabled:cursor-not-allowed"
                 >
-                  <span className="material-symbols-outlined text-[14px] mr-1">play_circle</span>
+                  <span className="material-symbols-outlined text-[16px]">play_circle</span>
                   Start DNS
-                </Button>
+                </button>
+              )}
+
+              {/* Warning below button */}
+              {warning && (
+                <div className="flex items-center gap-2 px-2 py-1.5 rounded text-xs bg-amber-500/10 text-amber-600 border border-amber-500/20">
+                  <span className="material-symbols-outlined text-[14px]">warning</span>
+                  <span>{warning}</span>
+                </div>
               )}
             </div>
           </div>
@@ -288,14 +266,14 @@ export default function MitmToolCard({
               onChange={(e) => setSudoPassword(e.target.value)}
               onKeyDown={(e) => { if (e.key === "Enter" && !loading) handleConfirmPassword(); }}
             />
-            {message && (
+            {modalError && (
               <div className="flex items-center gap-2 px-2 py-1.5 rounded text-xs bg-red-500/10 text-red-600">
                 <span className="material-symbols-outlined text-[14px]">error</span>
-                <span>{message.text}</span>
+                <span>{modalError}</span>
               </div>
             )}
             <div className="flex items-center justify-end gap-2">
-              <Button variant="ghost" size="sm" onClick={() => { setShowPasswordModal(false); setSudoPassword(""); setMessage(null); }} disabled={loading}>
+              <Button variant="ghost" size="sm" onClick={() => { setShowPasswordModal(false); setSudoPassword(""); setModalError(null); }} disabled={loading}>
                 Cancel
               </Button>
               <Button variant="primary" size="sm" onClick={handleConfirmPassword} loading={loading}>
diff --git a/src/app/api/cli-tools/antigravity-mitm/route.js b/src/app/api/cli-tools/antigravity-mitm/route.js
index 55eaa162..45dbcf95 100644
--- a/src/app/api/cli-tools/antigravity-mitm/route.js
+++ b/src/app/api/cli-tools/antigravity-mitm/route.js
@@ -43,7 +43,7 @@ export async function GET() {
       certExists: status.certExists || false,
       certTrusted: status.certTrusted || false,
       dnsStatus: status.dnsStatus || {},
-      hasCachedPassword: !!getCachedPassword(),
+      hasCachedPassword: !!getCachedPassword() || !!(await loadEncryptedPassword()),
       isAdmin: checkIsAdmin(),
     });
   } catch (error) {
diff --git a/src/app/api/tunnel/enable/route.js b/src/app/api/tunnel/enable/route.js
index 1b2fdc4a..2a2e08d3 100644
--- a/src/app/api/tunnel/enable/route.js
+++ b/src/app/api/tunnel/enable/route.js
@@ -1,9 +1,13 @@
 import { NextResponse } from "next/server";
 import { enableTunnel } from "@/lib/tunnel/tunnelManager";
 
+const DNS_WARMUP_DELAY_MS = 8000;
+
 export async function POST() {
   try {
     const result = await enableTunnel();
+    // Wait for DNS warmup to propagate at Cloudflare edge after tunnel registered
+    await new Promise((r) => setTimeout(r, DNS_WARMUP_DELAY_MS));
     return NextResponse.json(result);
   } catch (error) {
     console.error("Tunnel enable error:", error);
diff --git a/src/lib/tunnel/cloudflared.js b/src/lib/tunnel/cloudflared.js
index 61ae3e51..7a639c15 100644
--- a/src/lib/tunnel/cloudflared.js
+++ b/src/lib/tunnel/cloudflared.js
@@ -241,15 +241,26 @@ export async function spawnQuickTunnel(localPort, onUrlUpdate) {
       reject(new Error("Quick tunnel timed out"));
     }, 90000);
 
+    let lastUrl = null;
+
     const handleLog = (data) => {
       const msg = data.toString();
       const tunnelUrl = getQuickTunnelUrlFromLog(msg);
-      if (tunnelUrl && !resolved) {
+      if (!tunnelUrl) return;
+
+      if (!resolved) {
+        // First URL — resolve the promise, do NOT call onUrlUpdate (caller handles initial register)
         resolved = true;
+        lastUrl = tunnelUrl;
         clearTimeout(timeout);
         cleanup();
         resolve({ child, tunnelUrl });
-        // Notify caller of URL (for re-registration on URL change)
+        return;
+      }
+
+      // URL changed after initial connect — notify caller to re-register
+      if (tunnelUrl !== lastUrl) {
+        lastUrl = tunnelUrl;
         if (onUrlUpdate) onUrlUpdate(tunnelUrl);
       }
     };
diff --git a/src/lib/tunnel/tunnelManager.js b/src/lib/tunnel/tunnelManager.js
index 0f9c25b6..58ce87ca 100644
--- a/src/lib/tunnel/tunnelManager.js
+++ b/src/lib/tunnel/tunnelManager.js
@@ -12,6 +12,16 @@ const MAX_RECONNECT_ATTEMPTS = RECONNECT_DELAYS_MS.length;
 
 let isReconnecting = false;
 let exitHandlerRegistered = false;
+let reconnectTimeoutId = null;
+let manualDisabled = false;
+
+export function isTunnelManuallyDisabled() {
+  return manualDisabled;
+}
+
+export function isTunnelReconnecting() {
+  return isReconnecting;
+}
 
 function generateShortId() {
   let result = "";
@@ -43,6 +53,7 @@ async function registerTunnelUrl(shortId, tunnelUrl) {
 }
 
 export async function enableTunnel(localPort = 20128) {
+  manualDisabled = false;
   if (isCloudflaredRunning()) {
     const existing = loadState();
     if (existing?.tunnelUrl) {
@@ -56,15 +67,18 @@ export async function enableTunnel(localPort = 20128) {
   const existing = loadState();
   const shortId = existing?.shortId || generateShortId();
 
-  // Spawn quick tunnel, parse URL from cloudflared output
-  const { tunnelUrl } = await spawnQuickTunnel(localPort, async (url) => {
-    // Called on URL change (restart) - re-register new URL
+  // onUrlUpdate: only called when URL changes AFTER initial connect (not on first resolve)
+  const onUrlUpdate = async (url) => {
+    if (manualDisabled) return;
     await registerTunnelUrl(shortId, url);
     saveState({ shortId, machineId, tunnelUrl: url });
     await updateSettings({ tunnelEnabled: true, tunnelUrl: url });
-  });
+  };
 
-  // Register initial URL
+  // Spawn quick tunnel — resolve returns initial URL, onUrlUpdate handles subsequent changes
+  const { tunnelUrl } = await spawnQuickTunnel(localPort, onUrlUpdate);
+
+  // Register initial URL (exactly once)
   await registerTunnelUrl(shortId, tunnelUrl);
   saveState({ shortId, machineId, tunnelUrl });
   await updateSettings({ tunnelEnabled: true, tunnelUrl });
@@ -82,15 +96,19 @@ export async function enableTunnel(localPort = 20128) {
 }
 
 async function scheduleReconnect(attempt) {
-  if (isReconnecting) return;
+  if (isReconnecting || manualDisabled) return;
   isReconnecting = true;
 
   const delay = RECONNECT_DELAYS_MS[Math.min(attempt, RECONNECT_DELAYS_MS.length - 1)];
   console.log(`[Tunnel] Reconnecting in ${delay / 1000}s (attempt ${attempt + 1})...`);
 
-  await new Promise((r) => setTimeout(r, delay));
+  await new Promise((r) => { reconnectTimeoutId = setTimeout(r, delay); });
 
   try {
+    if (manualDisabled) {
+      isReconnecting = false;
+      return;
+    }
     const settings = await getSettings();
     if (!settings.tunnelEnabled) {
       isReconnecting = false;
@@ -109,8 +127,17 @@ async function scheduleReconnect(attempt) {
 }
 
 export async function disableTunnel() {
+  // Block any reconnect attempts before killing the process
+  manualDisabled = true;
+  isReconnecting = true;
+  if (reconnectTimeoutId) {
+    clearTimeout(reconnectTimeoutId);
+    reconnectTimeoutId = null;
+  }
+  setUnexpectedExitHandler(null);
+  exitHandlerRegistered = false;
+
   killCloudflared();
-  exitHandlerRegistered = false; // Reset handler flag when tunnel disabled
 
   const state = loadState();
   if (state) {
@@ -119,6 +146,9 @@ export async function disableTunnel() {
 
   await updateSettings({ tunnelEnabled: false, tunnelUrl: "" });
 
+  // Unblock reconnect lock — manualDisabled stays true to block Watchdog/NetworkMonitor
+  isReconnecting = false;
+
   return { success: true };
 }
 
diff --git a/src/mitm/manager.js b/src/mitm/manager.js
index 571099f0..1d96b73f 100644
--- a/src/mitm/manager.js
+++ b/src/mitm/manager.js
@@ -146,6 +146,15 @@ async function saveMitmSettings(enabled, password) {
   }
 }
 
+async function clearEncryptedPassword() {
+  if (!_updateSettings) return;
+  try {
+    await _updateSettings({ mitmSudoEncrypted: null });
+  } catch (e) {
+    err(`Failed to clear encrypted password: ${e.message}`);
+  }
+}
+
 async function loadEncryptedPassword() {
   if (!_getSettings) return null;
   try {
@@ -467,6 +476,7 @@ async function startServer(apiKey, sudoPassword) {
       // Detect wrong/missing password — clear cache and stop retry loop
       if (!IS_WIN && (msg.includes("incorrect password") || msg.includes("no password was provided"))) {
         setCachedPassword(null);
+        clearEncryptedPassword();
         mitmIsRestarting = true; // prevent scheduleMitmRestart from firing
       }
     });
@@ -603,5 +613,6 @@ module.exports = {
   getCachedPassword,
   setCachedPassword,
   loadEncryptedPassword,
+  clearEncryptedPassword,
   initDbHooks,
 };
diff --git a/src/shared/services/initializeApp.js b/src/shared/services/initializeApp.js
index f315ac46..42c665fa 100644
--- a/src/shared/services/initializeApp.js
+++ b/src/shared/services/initializeApp.js
@@ -1,5 +1,5 @@
 import { cleanupProviderConnections, getSettings, updateSettings, getApiKeys } from "@/lib/localDb";
-import { enableTunnel } from "@/lib/tunnel/tunnelManager";
+import { enableTunnel, isTunnelManuallyDisabled, isTunnelReconnecting } from "@/lib/tunnel/tunnelManager";
 import { killCloudflared, isCloudflaredRunning, ensureCloudflared } from "@/lib/tunnel/cloudflared";
 import { getMitmStatus, startMitm, loadEncryptedPassword, initDbHooks } from "@/mitm/manager";
 import { fileURLToPath } from "url";
@@ -135,12 +135,20 @@ function startWatchdog() {
   if (g.watchdogInterval) return;
   g.watchdogInterval = setInterval(async () => {
     try {
+      if (isTunnelManuallyDisabled()) return;
+      if (isTunnelReconnecting()) return;
+      if (g.tunnelRestartInProgress) return;
       const settings = await getSettings();
       if (!settings.tunnelEnabled) return;
       if (isCloudflaredRunning()) return;
       console.log("[Watchdog] Tunnel process is down, attempting recovery...");
-      await enableTunnel();
-      console.log("[Watchdog] Tunnel recovered");
+      g.tunnelRestartInProgress = true;
+      try {
+        await enableTunnel();
+        console.log("[Watchdog] Tunnel recovered");
+      } finally {
+        g.tunnelRestartInProgress = false;
+      }
     } catch (err) {
       console.log("[Watchdog] Recovery failed:", err.message);
     }
@@ -173,6 +181,7 @@ function startNetworkMonitor() {
 
   g.networkMonitorInterval = setInterval(async () => {
     try {
+      if (isTunnelManuallyDisabled()) return;
       const settings = await getSettings();
       if (!settings.tunnelEnabled) return;
 
@@ -190,6 +199,7 @@ function startNetworkMonitor() {
 
       // Skip if restart already in progress or restarted recently
       if (g.tunnelRestartInProgress) return;
+      if (isTunnelReconnecting()) return;
       if (now - g.lastTunnelRestartAt < NETWORK_RESTART_COOLDOWN_MS) return;
 
       const reason = wasSleep && networkChanged ? "sleep/wake + network change"