From cb9863b6bbc21035fcd05de230b96c1f3d2ab4e6 Mon Sep 17 00:00:00 2001
From: t8y2 <1156263951@qq.com>
Date: Sun, 17 May 2026 17:06:59 +0800
Subject: [PATCH] feat(drivers): prompt for agent driver updates

---
 src/App.vue                           | 112 ++++++++++++++++++++++++++
 src/lib/agentDriverUpdatePrompt.ts    |  50 ++++++++++++
 src/stores/connectionStore.ts         |  10 +++
 tests/agentDriverUpdatePrompt.test.ts |  75 +++++++++++++++++
 4 files changed, 247 insertions(+)
 create mode 100644 src/lib/agentDriverUpdatePrompt.ts
 create mode 100644 tests/agentDriverUpdatePrompt.test.ts

diff --git a/src/App.vue b/src/App.vue
index 2fe103af1..e7fa58c28 100644
--- a/src/App.vue
+++ b/src/App.vue
@@ -1,6 +1,7 @@
 <script setup lang="ts">
 import { ref, computed, watch, onMounted, onUnmounted, nextTick } from "vue";
 import { useI18n } from "vue-i18n";
+import { invoke } from "@tauri-apps/api/core";
 import { TooltipProvider } from "@/components/ui/tooltip";
 import AiAssistant from "@/components/editor/AiAssistant.vue";
 import QueryHistory from "@/components/editor/QueryHistory.vue";
@@ -48,10 +49,17 @@ import { isPreviewTab } from "@/lib/tabPresentation";
 import { supportsSqlFileExecution } from "@/lib/databaseCapabilities";
 import { classifyAiSqlExecution } from "@/lib/aiSqlExecutionPolicy";
 import { buildHistoryAiAnalysisPrompt } from "@/lib/historyAiAnalysis";
+import {
+  agentDriverUpdateIgnoreKey,
+  findAgentDriverUpdatePrompt,
+  type AgentDriverUpdatePrompt,
+  type AgentDriverUpdateState,
+} from "@/lib/agentDriverUpdatePrompt";
 import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog";
 import { Button } from "@/components/ui/button";
 import { Input } from "@/components/ui/input";
 import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
+import type { ConnectionConfig } from "@/types/database";
 import type { HistoryEntry } from "@/lib/tauri";
 
 const { t } = useI18n();
@@ -80,6 +88,8 @@ const { setupFileDrop } = useFileDrop();
 
 const isDesktop = isTauriRuntime();
 const UPDATE_CHECK_INTERVAL_MS = 60 * 60 * 1000;
+const AGENT_DRIVER_UPDATE_IGNORES_STORAGE_KEY = "dbx-agent-driver-update-ignores";
+type AgentDriverUpdateChoice = "update" | "skip" | "ignore";
 let updateCheckTimer: ReturnType<typeof setInterval> | undefined;
 const needsAuth = ref(!isDesktop);
 const authenticated = ref(isDesktop);
@@ -88,6 +98,9 @@ const setupRequired = ref(false);
 const showConnectionDialog = ref(false);
 const showSettingsDialog = ref(false);
 const showDriverStore = ref(false);
+const showAgentDriverUpdateDialog = ref(false);
+const agentDriverUpdatePrompt = ref<AgentDriverUpdatePrompt | null>(null);
+const isInstallingAgentDriverUpdate = ref(false);
 const showHistory = ref(false);
 const showAiPanel = ref(localStorage.getItem("dbx-ai-panel-open") !== "false");
 const { sidebarWidth, aiPanelWidth, historyWidth, startSidebarResize, startAiPanelResize, startHistoryResize } =
@@ -112,6 +125,70 @@ const activeConnection = computed(() => {
   return tab ? connectionStore.getConfig(tab.connectionId) : undefined;
 });
 
+let resolveAgentDriverUpdateChoice: ((choice: AgentDriverUpdateChoice) => void) | null = null;
+
+function loadIgnoredAgentDriverUpdates(): Set<string> {
+  try {
+    const value = JSON.parse(localStorage.getItem(AGENT_DRIVER_UPDATE_IGNORES_STORAGE_KEY) || "[]");
+    return new Set(Array.isArray(value) ? value.filter((item): item is string => typeof item === "string") : []);
+  } catch {
+    return new Set();
+  }
+}
+
+function saveIgnoredAgentDriverUpdate(prompt: AgentDriverUpdatePrompt) {
+  const ignored = loadIgnoredAgentDriverUpdates();
+  ignored.add(agentDriverUpdateIgnoreKey(prompt.dbType, prompt.latestVersion));
+  localStorage.setItem(AGENT_DRIVER_UPDATE_IGNORES_STORAGE_KEY, JSON.stringify([...ignored]));
+}
+
+function requestAgentDriverUpdateChoice(prompt: AgentDriverUpdatePrompt): Promise<AgentDriverUpdateChoice> {
+  agentDriverUpdatePrompt.value = prompt;
+  showAgentDriverUpdateDialog.value = true;
+  return new Promise((resolve) => {
+    resolveAgentDriverUpdateChoice = resolve;
+  });
+}
+
+function chooseAgentDriverUpdate(choice: AgentDriverUpdateChoice) {
+  showAgentDriverUpdateDialog.value = false;
+  resolveAgentDriverUpdateChoice?.(choice);
+  resolveAgentDriverUpdateChoice = null;
+}
+
+function onAgentDriverUpdateOpenChange(value: boolean) {
+  if (value) {
+    showAgentDriverUpdateDialog.value = true;
+    return;
+  }
+  if (isInstallingAgentDriverUpdate.value) return;
+  chooseAgentDriverUpdate("skip");
+}
+
+async function checkAgentDriverUpdateBeforeConnect(config: ConnectionConfig) {
+  if (!isDesktop) return;
+  const drivers = await invoke<AgentDriverUpdateState[]>("list_installed_agents").catch(() => []);
+  const prompt = findAgentDriverUpdatePrompt(config, drivers, loadIgnoredAgentDriverUpdates());
+  if (!prompt) return;
+
+  const choice = await requestAgentDriverUpdateChoice(prompt);
+  if (choice === "ignore") {
+    saveIgnoredAgentDriverUpdate(prompt);
+    return;
+  }
+  if (choice === "skip") return;
+
+  isInstallingAgentDriverUpdate.value = true;
+  try {
+    await invoke("install_agent", { dbType: prompt.dbType });
+    toast(`${prompt.label} 驱动已更新到 v${prompt.latestVersion}`);
+  } finally {
+    isInstallingAgentDriverUpdate.value = false;
+  }
+}
+
+connectionStore.setBeforeConnectHandler(checkAgentDriverUpdateBeforeConnect);
+
 function restoreHistorySql(sql: string, entry: HistoryEntry) {
   const tab = activeTab.value;
   if (tab?.mode === "query") {
@@ -648,6 +725,7 @@ onMounted(async () => {
 });
 
 onUnmounted(() => {
+  connectionStore.setBeforeConnectHandler(null);
   cleanupTauriListeners();
   if (updateCheckTimer) {
     clearInterval(updateCheckTimer);
@@ -881,6 +959,40 @@ onUnmounted(() => {
           @download-and-install="downloadAndInstallUpdate"
           @restart="restartApp"
         />
+        <Dialog :open="showAgentDriverUpdateDialog" @update:open="onAgentDriverUpdateOpenChange">
+          <DialogContent class="sm:max-w-[440px]">
+            <DialogHeader>
+              <DialogTitle>驱动有新版本</DialogTitle>
+            </DialogHeader>
+            <div v-if="agentDriverUpdatePrompt" class="space-y-3 text-sm text-muted-foreground">
+              <p>
+                {{ agentDriverUpdatePrompt.label }} 驱动可从 v{{ agentDriverUpdatePrompt.currentVersion }} 更新到 v{{
+                  agentDriverUpdatePrompt.latestVersion
+                }}。
+              </p>
+              <p>可以先更新驱动再连接，也可以本次直接连接。</p>
+            </div>
+            <DialogFooter class="gap-2 sm:gap-2">
+              <Button
+                variant="ghost"
+                :disabled="isInstallingAgentDriverUpdate"
+                @click="chooseAgentDriverUpdate('ignore')"
+              >
+                本版本不再提醒
+              </Button>
+              <Button
+                variant="outline"
+                :disabled="isInstallingAgentDriverUpdate"
+                @click="chooseAgentDriverUpdate('skip')"
+              >
+                直接连接
+              </Button>
+              <Button :disabled="isInstallingAgentDriverUpdate" @click="chooseAgentDriverUpdate('update')">
+                {{ isInstallingAgentDriverUpdate ? "更新中..." : "更新后连接" }}
+              </Button>
+            </DialogFooter>
+          </DialogContent>
+        </Dialog>
 
         <Transition name="toast">
           <div
diff --git a/src/lib/agentDriverUpdatePrompt.ts b/src/lib/agentDriverUpdatePrompt.ts
new file mode 100644
index 000000000..19f276d17
--- /dev/null
+++ b/src/lib/agentDriverUpdatePrompt.ts
@@ -0,0 +1,50 @@
+import type { ConnectionConfig, DatabaseType } from "@/types/database";
+import { supportsDriverManagement } from "@/lib/databaseCapabilities";
+
+export interface AgentDriverUpdateState {
+  db_type: string;
+  label: string;
+  installed: boolean;
+  installed_version: string | null;
+  version: string;
+  update_available: boolean;
+}
+
+export interface AgentDriverUpdatePrompt {
+  dbType: string;
+  label: string;
+  currentVersion: string;
+  latestVersion: string;
+}
+
+export function agentDriverUpdateKeyForConnection(
+  connection: Pick<ConnectionConfig, "db_type" | "driver_profile">,
+): string | null {
+  if (connection.db_type === "oracle" && connection.driver_profile === "oracle-10g") return "oracle-10g";
+  if (!supportsDriverManagement(connection.db_type as DatabaseType)) return null;
+  return connection.db_type;
+}
+
+export function agentDriverUpdateIgnoreKey(dbType: string, version: string): string {
+  return `${dbType}@${version}`;
+}
+
+export function findAgentDriverUpdatePrompt(
+  connection: Pick<ConnectionConfig, "db_type" | "driver_profile">,
+  drivers: readonly AgentDriverUpdateState[],
+  ignoredVersions: ReadonlySet<string> = new Set(),
+): AgentDriverUpdatePrompt | null {
+  const dbType = agentDriverUpdateKeyForConnection(connection);
+  if (!dbType) return null;
+
+  const driver = drivers.find((item) => item.db_type === dbType);
+  if (!driver?.installed || !driver.update_available || !driver.version || !driver.installed_version) return null;
+  if (ignoredVersions.has(agentDriverUpdateIgnoreKey(dbType, driver.version))) return null;
+
+  return {
+    dbType,
+    label: driver.label || dbType,
+    currentVersion: driver.installed_version,
+    latestVersion: driver.version,
+  };
+}
diff --git a/src/stores/connectionStore.ts b/src/stores/connectionStore.ts
index 1b97b7e54..6f47a17b1 100644
--- a/src/stores/connectionStore.ts
+++ b/src/stores/connectionStore.ts
@@ -47,6 +47,8 @@ interface PersistedTreeChildrenLoadResult {
   isStale: boolean;
 }
 
+type BeforeConnectHandler = (config: ConnectionConfig) => Promise<void>;
+
 function redisDbLabel(db: number, loadedKeyCount?: number, totalKeyCount?: number): string {
   if (totalKeyCount == null) return `db${db}`;
   return `db${db} (${loadedKeyCount ?? 0}/${totalKeyCount})`;
@@ -120,6 +122,7 @@ export const useConnectionStore = defineStore("connection", () => {
   const sidebarLayout = ref<SidebarLayout>(emptyLayout());
   let layoutPersistTimer: ReturnType<typeof setTimeout> | null = null;
   const staleTreeRefreshIds = new Set<string>();
+  let beforeConnectHandler: BeforeConnectHandler | null = null;
 
   function startEditing(id: string) {
     editingConnectionId.value = id;
@@ -566,6 +569,7 @@ export const useConnectionStore = defineStore("connection", () => {
     const pendingNode = findNode(treeNodes.value, config.id);
     if (pendingNode) pendingNode.isLoading = true;
     try {
+      await beforeConnectHandler?.(config);
       const id = await api.connectDb(config);
       activeConnectionId.value = id;
       connectedIds.value.add(id);
@@ -620,6 +624,7 @@ export const useConnectionStore = defineStore("connection", () => {
       throw error;
     }
     try {
+      await beforeConnectHandler?.(config);
       await api.connectDb(config);
       connectedIds.value.add(connectionId);
       activeConnectionId.value = connectionId;
@@ -630,6 +635,10 @@ export const useConnectionStore = defineStore("connection", () => {
     }
   }
 
+  function setBeforeConnectHandler(handler: BeforeConnectHandler | null) {
+    beforeConnectHandler = handler;
+  }
+
   async function loadDatabases(connectionId: string, options?: LoadTreeOptions) {
     const node = findNode(treeNodes.value, connectionId);
     if (!node) return;
@@ -1671,6 +1680,7 @@ export const useConnectionStore = defineStore("connection", () => {
     connect,
     disconnect,
     ensureConnected,
+    setBeforeConnectHandler,
     initFromDisk,
     loadDatabases,
     loadRedisDatabases,
diff --git a/tests/agentDriverUpdatePrompt.test.ts b/tests/agentDriverUpdatePrompt.test.ts
new file mode 100644
index 000000000..010db9c56
--- /dev/null
+++ b/tests/agentDriverUpdatePrompt.test.ts
@@ -0,0 +1,75 @@
+import assert from "node:assert/strict";
+import test from "node:test";
+import {
+  agentDriverUpdateIgnoreKey,
+  agentDriverUpdateKeyForConnection,
+  findAgentDriverUpdatePrompt,
+} from "../src/lib/agentDriverUpdatePrompt.ts";
+
+const oracleConnection = {
+  id: "conn",
+  name: "Oracle",
+  db_type: "oracle",
+  driver_profile: undefined,
+} as const;
+
+test("uses the driver profile for versioned agent drivers", () => {
+  assert.equal(agentDriverUpdateKeyForConnection({ db_type: "oracle", driver_profile: "oracle-10g" }), "oracle-10g");
+  assert.equal(agentDriverUpdateKeyForConnection({ db_type: "oracle" }), "oracle");
+  assert.equal(agentDriverUpdateKeyForConnection({ db_type: "postgres" }), null);
+});
+
+test("finds installed agent drivers with available updates", () => {
+  assert.deepEqual(
+    findAgentDriverUpdatePrompt(oracleConnection, [
+      {
+        db_type: "oracle",
+        label: "Oracle",
+        installed: true,
+        installed_version: "1.0.0",
+        version: "1.1.0",
+        update_available: true,
+      },
+    ]),
+    {
+      dbType: "oracle",
+      label: "Oracle",
+      currentVersion: "1.0.0",
+      latestVersion: "1.1.0",
+    },
+  );
+});
+
+test("does not prompt for missing, current, or ignored driver updates", () => {
+  assert.equal(
+    findAgentDriverUpdatePrompt(oracleConnection, [
+      {
+        db_type: "oracle",
+        label: "Oracle",
+        installed: false,
+        installed_version: null,
+        version: "1.1.0",
+        update_available: true,
+      },
+    ]),
+    null,
+  );
+
+  assert.equal(
+    findAgentDriverUpdatePrompt(
+      oracleConnection,
+      [
+        {
+          db_type: "oracle",
+          label: "Oracle",
+          installed: true,
+          installed_version: "1.0.0",
+          version: "1.1.0",
+          update_available: true,
+        },
+      ],
+      new Set([agentDriverUpdateIgnoreKey("oracle", "1.1.0")]),
+    ),
+    null,
+  );
+});