From f4dbfe0cc3cbd68b5e819b03b75c108c1d5bea8a Mon Sep 17 00:00:00 2001
From: zipg <geeky@foxmail.com>
Date: Mon, 17 Aug 2026 01:52:15 +0800
Subject: [PATCH] fix(query): deduplicate timeout setting actions

---
 .../src/components/common/QueryErrorActions.vue  |  7 +++++--
 .../common/__tests__/QueryErrorActions.spec.ts   | 16 ++++++++++++++++
 .../src/lib/__tests__/sql/queryError.spec.ts     | 12 ++++++++++++
 apps/desktop/src/lib/sql/queryError.ts           |  9 +++++++--
 4 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/apps/desktop/src/components/common/QueryErrorActions.vue b/apps/desktop/src/components/common/QueryErrorActions.vue
index 81cdb1c91a..edb9f6b381 100644
--- a/apps/desktop/src/components/common/QueryErrorActions.vue
+++ b/apps/desktop/src/components/common/QueryErrorActions.vue
@@ -1,4 +1,5 @@
 <script setup lang="ts">
+import { computed } from "vue";
 import { Bot, Wrench } from "@lucide/vue";
 import { useI18n } from "vue-i18n";
 import { Button } from "@/components/ui/button";
@@ -18,14 +19,16 @@ const emit = defineEmits<{
 }>();
 
 const { t } = useI18n();
+const showConnectionTimeout = computed(() => !!props.connectionId && isConnectionTimeoutErrorMessage(props.errorMessage, props.backendError));
+const showQueryTimeout = computed(() => !!props.connectionId && !showConnectionTimeout.value && isQueryTimeoutErrorMessage(props.errorMessage, props.backendError));
 </script>
 
 <template>
-  <Button v-if="connectionId && isConnectionTimeoutErrorMessage(errorMessage, backendError)" variant="outline" size="sm" class="h-7 gap-1.5 px-2.5 text-xs" @click="emit('changeConnectionTimeout')">
+  <Button v-if="showConnectionTimeout" variant="outline" size="sm" class="h-7 gap-1.5 px-2.5 text-xs" @click="emit('changeConnectionTimeout')">
     <Wrench class="h-3.5 w-3.5" />
     {{ t("editor.changeConnectionTimeout") }}
   </Button>
-  <Button v-if="connectionId && isQueryTimeoutErrorMessage(errorMessage, backendError)" variant="outline" size="sm" class="h-7 gap-1.5 px-2.5 text-xs" @click="emit('changeQueryTimeout')">
+  <Button v-if="showQueryTimeout" variant="outline" size="sm" class="h-7 gap-1.5 px-2.5 text-xs" @click="emit('changeQueryTimeout')">
     <Wrench class="h-3.5 w-3.5" />
     {{ t("editor.changeQueryTimeout") }}
   </Button>
diff --git a/apps/desktop/src/components/common/__tests__/QueryErrorActions.spec.ts b/apps/desktop/src/components/common/__tests__/QueryErrorActions.spec.ts
index f66b930196..94d902eeaa 100644
--- a/apps/desktop/src/components/common/__tests__/QueryErrorActions.spec.ts
+++ b/apps/desktop/src/components/common/__tests__/QueryErrorActions.spec.ts
@@ -41,4 +41,20 @@ describe("QueryErrorActions", () => {
     expect(host.textContent).not.toContain("Change connection timeout");
     expect(host.textContent).toContain("Change query timeout");
   });
+
+  it("shows only connection timeout settings for connect-stage agent RPC timeouts", async () => {
+    const { host } = await mountActions("Agent RPC call timed out at connect");
+    const labels = Array.from(host.querySelectorAll("button"), (button) => button.textContent?.trim());
+
+    expect(labels).toContain("Change connection timeout");
+    expect(labels).not.toContain("Change query timeout");
+  });
+
+  it("does not offer query timeout settings for non-query Agent RPC stages", async () => {
+    for (const stage of ["request", "validate", "cancel", "close"]) {
+      const { host } = await mountActions(`Agent RPC call timed out at ${stage}`);
+
+      expect(host.textContent).not.toContain("Change query timeout");
+    }
+  });
 });
diff --git a/apps/desktop/src/lib/__tests__/sql/queryError.spec.ts b/apps/desktop/src/lib/__tests__/sql/queryError.spec.ts
index aa8d6fad39..46bc3a856f 100644
--- a/apps/desktop/src/lib/__tests__/sql/queryError.spec.ts
+++ b/apps/desktop/src/lib/__tests__/sql/queryError.spec.ts
@@ -64,6 +64,18 @@ describe("isQueryTimeoutErrorMessage", () => {
     expect(isQueryTimeoutErrorMessage("Agent RPC error (-1): Agent RPC call timed out (120s)")).toBe(true);
   });
 
+  it("classifies only execute and fetch Agent RPC stages as query timeouts", () => {
+    for (const stage of ["execute", "fetch"]) {
+      expect(isQueryTimeoutErrorMessage(`Agent RPC call timed out at ${stage}`)).toBe(true);
+    }
+
+    for (const stage of ["request", "checkout", "connect", "validate", "cancel", "close"]) {
+      expect(isQueryTimeoutErrorMessage(`Agent RPC call timed out at ${stage}`)).toBe(false);
+    }
+
+    expect(isConnectionTimeoutErrorMessage("Agent RPC call timed out at connect")).toBe(true);
+  });
+
   it("detects structured query operation timeouts before localized text matching", () => {
     const timeoutError = {
       version: 1,
diff --git a/apps/desktop/src/lib/sql/queryError.ts b/apps/desktop/src/lib/sql/queryError.ts
index 7602a33601..0e49ed633c 100644
--- a/apps/desktop/src/lib/sql/queryError.ts
+++ b/apps/desktop/src/lib/sql/queryError.ts
@@ -26,8 +26,13 @@ export function isQueryTimeoutErrorMessage(message: string, backendError?: Backe
   // Agent RPC client-side timeout (tokio::time::timeout in agent_driver.rs). This is the
   // fallback when JDBC setQueryTimeout never fires (unsupported/unresponsive driver), so the
   // backend already treats it as a query timeout (is_agent_rpc_timeout_error in query.rs) —
-  // surface the same action here to stay consistent with the backend.
-  if (/\bagent rpc call timed out\b/.test(lower)) return true;
+  // surface the same action here to stay consistent with the backend. Connection and other
+  // infrastructure stages must be excluded first because their messages can share this prefix.
+  const agentTimeout = /\bagent rpc call timed out(?: at ([a-z_]+))?\b/.exec(lower);
+  if (agentTimeout) {
+    const stage = agentTimeout[1];
+    return stage === undefined || QUERY_TIMEOUT_STAGES.has(stage);
+  }
   if (/\b(?:canceling|cancelling|canceled|cancelled)\b[\s\S]{0,80}\bstatement\b[\s\S]{0,80}\btimeout\b/.test(lower)) return true;
   if (/\b(?:connection|connect|pool|checkout|metadata|loading|health check|cancel request|ssh|tunnel)\b/.test(lower)) return false;
   return (