feat(ai): add connection tree context

apps/desktop/src/App.vue

@@ -62,7 +62,7 @@ import { isTauriRuntime } from "@/lib/backend/tauriRuntime";
 import { openQueryResultArchiveFile } from "@/lib/query/queryResultArchiveFile";
 import { rememberExternalSqlFileTarget, resolveExternalSqlFileTarget, unassociatedExternalSqlFileTarget } from "@/lib/sql/externalSqlFileTarget";
 import { externalSqlFileOpenErrorMessage, readBrowserSqlFile, sqlFileTitleFromPath } from "@/lib/sql/sqlFileOpen";
-import type { ConnectionConfig, DatabaseType, ObjectSourceKind, QueryTab } from "@/types/database";
+import type { ConnectionConfig, DatabaseType, ObjectSourceKind, QueryTab, TreeNode } from "@/types/database";
 import { parseConnectionDeepLink, type ConnectionDeepLinkDraft } from "@/lib/connection/connectionDeepLink";
 import {
   isBrowserReloadShortcut,
@@ -109,6 +109,7 @@ import { savedSqlErrorMessage } from "@/lib/savedSql/savedSqlErrors";
 import { savedSqlDefaultTargetForWrite } from "@/lib/savedSql/savedSqlExecutionTarget";
 import { countActiveUpdateBlockingTasks } from "@/lib/app/appUpdateTaskGuard";
 import { initSavedSqlEditorPositions } from "@/lib/app/savedSqlEditorPosition";
+import { hasTreeNodeDatabaseContext } from "@/lib/sidebar/treeNodeContext";
 import { isSchemaAware, isSingleDatabase, usesTreeSchemaMode } from "@/lib/database/databaseFeatureSupport";
 import { codeMirrorSqlDialect, connectionUsesDatabaseObjectTreeMode, effectiveDatabaseTypeForConnection } from "@/lib/database/jdbcDialect";
 import { sqlFormatDialectForDbType } from "@/lib/sql/sqlFormatter";
@@ -120,7 +121,7 @@ import { Button } from "@/components/ui/button";
 import { Input } from "@/components/ui/input";
 import { SearchableSelect } from "@/components/ui/searchable-select";
 import type { HistoryEntry } from "@/lib/backend/tauri";
-import type { AiAction } from "@/lib/ai/ai";
+import { resolveDefaultAiSchema, type AiAction } from "@/lib/ai/ai";
 import ExternalSqlFileChangeDialog from "@/components/editor/ExternalSqlFileChangeDialog.vue";
 
 const AiAssistant = defineAsyncComponent(() => import("@/components/editor/AiAssistant.vue"));
@@ -139,6 +140,8 @@ const QueryEditorObjectSourceDialog = defineAsyncComponent(() => import("@/compo
 type AiAssistantHandle = {
   triggerAction: (action: AiAction, instruction?: string) => void;
   setPrompt: (text: string) => void;
+  addTableMention: (target: { schema?: string; table: string }) => void;
+  clearContextReferences: () => void;
 };
 
 const { t } = useI18n();
@@ -857,6 +860,58 @@ function sendSelectionToAi(sql: string) {
   invokeWhenAiReady((handle) => handle.setPrompt(sql));
 }
 
+let addToAiRequestId = 0;
+
+async function addToAi(node: TreeNode) {
+  if ((node.type !== "connection" && node.type !== "database" && node.type !== "table") || !node.connectionId) return;
+  const connection = connectionStore.getConfig(node.connectionId);
+  if (!connection) return;
+  const requestId = ++addToAiRequestId;
+
+  try {
+    await connectionStore.ensureConnected(node.connectionId);
+    if (requestId !== addToAiRequestId) return;
+    connectionStore.activeConnectionId = node.connectionId;
+
+    let target: { database: string; schema?: string; catalog?: string } | null = null;
+    if (node.type === "connection") {
+      const options = await getDatabaseOptions(node.connectionId);
+      if (requestId !== addToAiRequestId) return;
+      target =
+        connection.db_type === "dameng"
+          ? {
+              database: resolveDefaultDatabase(connection, []),
+              schema: resolveDefaultAiSchema(connection, options),
+            }
+          : {
+              database: resolveDefaultDatabase(connection, options),
+              schema: connection.default_schema,
+            };
+    } else if (hasTreeNodeDatabaseContext(node)) {
+      target = { database: node.database, schema: node.schema, catalog: node.catalog };
+    }
+    if (!target) return;
+
+    const currentTab = activeTab.value;
+    const contextChanged = currentTab?.connectionId !== node.connectionId || currentTab?.database !== target.database || (currentTab?.schema || "") !== (target.schema || "") || (currentTab?.catalog || "") !== (target.catalog || "");
+
+    const contextTab = queryStore.tabs.find((tab) => tab.connectionId === node.connectionId && tab.database === target.database && (tab.schema || "") === (target.schema || "") && (tab.catalog || "") === (target.catalog || ""));
+    if (contextTab) {
+      queryStore.switchTab(contextTab.id);
+    } else {
+      queryStore.createTab(node.connectionId, target.database, undefined, "query", target.schema, undefined, target.catalog);
+    }
+
+    openRightSidebarPanel("ai");
+    invokeWhenAiReady((handle) => {
+      if (contextChanged) handle.clearContextReferences();
+      if (node.type === "table") handle.addTableMention({ schema: node.schema, table: node.label });
+    });
+  } catch (e: any) {
+    toast(t("connection.connectFailed", { message: translateBackendError(t, e) }), 5000);
+  }
+}
+
 function openAiPanel() {
   openRightSidebarPanel("ai");
 }
@@ -2669,6 +2724,7 @@ onUnmounted(() => {
             @start-resize="startSidebarResize"
             @collapse="setSidebarOpen(false)"
             @open-settings="(initialTab) => openSettings(initialTab ?? 'appearance')"
+            @add-to-ai="addToAi"
           />
           <div v-show="!sidebarOpen" class="flex h-full w-8 shrink-0 items-start justify-center border-r bg-background/80 pt-2" :class="isClassicLayout ? '' : 'rounded-md border border-border/80'">
             <Button variant="ghost" size="icon" class="h-7 w-7" :title="t('sidebar.expand')" :aria-label="t('sidebar.expand')" @click="setSidebarOpen(true)">

apps/desktop/src/components/editor/AiAssistant.vue

@@ -1063,6 +1063,8 @@ async function loadDatabases(connection = props.connection): Promise<string[]> {
 async function changeConnection(connectionId: string) {
   const conn = connectionStore.getConfig(connectionId);
   if (!conn) return;
+  if (props.connection?.id === connectionId) return;
+  clearContextReferences();
   connectionStore.activeConnectionId = connectionId;
   const tab = props.tab;
   const tabId = tab ? tab.id : queryStore.createTab(connectionId, resolveDefaultDatabase(conn, []));
@@ -1087,6 +1089,8 @@ function changeNamespace(value: string) {
   const connection = props.connection;
   if (!tab || !connection) return;
   const namespace = decodeSelectableDatabaseValue(connection.db_type, value);
+  if (resolveAiNamespaceSelection(tab, connection).value === namespace) return;
+  clearContextReferences();
   if (resolveAiNamespaceSelection(tab, connection).kind === "schema") {
     queryStore.updateSchema(tab.id, namespace || undefined);
   } else {
@@ -2942,7 +2946,23 @@ function setPrompt(text: string) {
   nextTick(() => promptTextareaRef.value?.focus());
 }
 
-defineExpose({ triggerAction, setPrompt });
+function addTableMention(target: { schema?: string; table: string }) {
+  const table = target.table.trim();
+  if (!table) return;
+  addSelectedMention({ kind: "table", schema: target.schema, name: table, tableType: "TABLE" });
+  nextTick(() => promptTextareaRef.value?.focus());
+}
+
+function clearContextReferences() {
+  selectedMentions.value = [];
+  selectedSqlFileMentions.value = [];
+  mentionCache.value = {};
+  mentionCandidates.value = [];
+  mentionOpen.value = false;
+  mentionError.value = "";
+}
+
+defineExpose({ triggerAction, setPrompt, addTableMention, clearContextReferences });
 
 const messageRenderer = computed(() => {
   const appearance = aiCodeAppearance.value;

apps/desktop/src/components/layout/AppSidebar.vue

@@ -12,6 +12,7 @@ import ConnectionTree from "@/components/sidebar/ConnectionTree.vue";
 import { applyConnectionMultiSelection, emptyConnectionMultiSelection, isExitConnectionMultiSelectionShortcut } from "@/lib/sidebar/sidebarConnectionMultiSelect";
 import { useConnectionStore } from "@/stores/connectionStore";
 import { useToast } from "@/composables/useToast";
+import type { TreeNode } from "@/types/database";
 
 defineProps<{
   sidebarWidth: number;
@@ -24,6 +25,7 @@ const emit = defineEmits<{
   startResize: [event: MouseEvent];
   collapse: [];
   "open-settings": [initialTab: string];
+  "add-to-ai": [node: TreeNode];
 }>();
 
 type ImportSource = "dbx" | "navicat" | "dbeaver" | "datagrip";
@@ -264,7 +266,7 @@ defineExpose({ focusSearch });
         </template>
       </div>
       <div class="flex-1 min-h-0">
-        <ConnectionTree ref="connectionTreeRef" @open-settings="(initialTab) => emit('open-settings', initialTab)" />
+        <ConnectionTree ref="connectionTreeRef" @open-settings="(initialTab) => emit('open-settings', initialTab)" @add-to-ai="(node) => emit('add-to-ai', node)" />
       </div>
     </div>
     <div class="panel-resize-handle panel-resize-handle--right" @mousedown="emit('startResize', $event)" />

apps/desktop/src/components/sidebar/ConnectionTree.vue

@@ -97,6 +97,7 @@ const sidebarContextMenuRef = ref<{ close: () => void } | null>(null);
 const sidebarContextMenuItems = ref<ContextMenuItem[]>([]);
 const emit = defineEmits<{
   "open-settings": [initialTab: string];
+  "add-to-ai": [node: TreeNode];
 }>();
 
 const sidebarContextMenuTarget = ref<SidebarActionTarget | null>(null);
@@ -2159,6 +2160,7 @@ defineExpose({ focusSearch, createNewGroup, collapseAllTreeNodes });
       @open-visible-schemas="openSidebarVisibleSchemas"
       @open-visible-nacos-namespaces="openSidebarVisibleNacosNamespaces"
       @open-table-name-filters="openSidebarTableNameFilters"
+      @add-to-ai="(node) => emit('add-to-ai', node)"
       @request-group-rename="startRenamingCreatedGroup"
       @request-saved-sql-rename="startRenamingSavedSqlNode"
       @open-danger-dialog="openSidebarDangerDialog"

apps/desktop/src/components/sidebar/SidebarTreeRuntimeHost.vue

@@ -57,6 +57,7 @@ import {
   X,
   Settings2,
   GitBranch,
+  Sparkles,
 } from "@lucide/vue";
 import type { ContextMenuItem } from "@/components/ui/CustomContextMenu.vue";
 import { CONNECTION_ATTEMPT_CANCELLED_MESSAGE, useConnectionStore } from "@/stores/connectionStore";
@@ -92,6 +93,7 @@ import {
   supportsDatabaseSearch,
   supportsConnectionDatabaseBrowser,
   supportsConnectionQueryActions,
+  supportsAiAssistantContext,
   supportsFieldLineage,
   supportsObjectBrowserTreeNode,
   supportsSchemaDiagram,
@@ -374,6 +376,7 @@ const emit = defineEmits<{
   "open-visible-schemas": [node: TreeNode];
   "open-visible-nacos-namespaces": [node: TreeNode];
   "open-table-name-filters": [node: TreeNode];
+  "add-to-ai": [node: TreeNode];
   "open-danger-dialog": [request: SidebarDangerDialogRequest];
   "open-dialog-controller": [controller: Record<string, any> | null];
   "open-install-extension": [node: TreeNode];
@@ -4581,8 +4584,12 @@ function buildConnectionSidebarMenu(context: SidebarMenuFactoryContext): boolean
     items.push({ label: t("contextMenu.copyName"), action: copyName, icon: Copy, shortcut: shortcutCopyName.value });
     items.push({ label: "", separator: true });
     const supportsQueryActions = supportsConnectionQueryActions(currentDatabaseType());
+    const supportsAiContext = supportsAiAssistantContext(currentDatabaseType());
     if (supportsQueryActions) {
       items.push({ label: t("contextMenu.newQuery"), action: newQuery, icon: TerminalSquare });
+      if (supportsAiContext) {
+        items.push({ label: t("contextMenu.addToAi"), action: () => emit("add-to-ai", node), icon: Sparkles });
+      }
     }
     const connectionWorkspace = node.connectionId ? driverProfileDatabaseWorkspace(connectionStore.getConfig(node.connectionId)?.driver_profile) : undefined;
     if (connectionWorkspace?.entryScopes.includes("connection")) {
@@ -4777,6 +4784,9 @@ function buildDatabaseSidebarMenu(context: SidebarMenuFactoryContext): boolean {
     }
     if (supportsConnectionQueryActions(currentDatabaseType())) {
       items.push({ label: t("contextMenu.newQuery"), action: newQuery, icon: TerminalSquare });
+      if (node.type === "database" && supportsAiAssistantContext(currentDatabaseType())) {
+        items.push({ label: t("contextMenu.addToAi"), action: () => emit("add-to-ai", node), icon: Sparkles });
+      }
       const sqlHistoryMenu = savedSqlHistorySubmenu();
       if (sqlHistoryMenu) items.push(sqlHistoryMenu);
     }
@@ -5072,6 +5082,9 @@ function buildObjectSidebarMenu(context: SidebarMenuFactoryContext): boolean {
       items.push({ label: t("contextMenu.viewData"), action: openDataImmediately, icon: TableProperties });
       items.push({ label: t("contextMenu.openInNewDataTab"), action: openDataInNewTabImmediately, icon: CopyPlus });
       items.push({ label: t("contextMenu.newQuery"), action: newQuery, icon: TerminalSquare });
+      if (supportsAiAssistantContext(currentDatabaseType())) {
+        items.push({ label: t("contextMenu.addToAi"), action: () => emit("add-to-ai", node), icon: Sparkles });
+      }
       items.push({ label: "", separator: true });
       items.push(exportDataSubmenu(false));
       items.push({ label: t("contextMenu.refreshChildren"), action: refresh, icon: RefreshCw, shortcut: shortcutRefresh });
@@ -5080,6 +5093,9 @@ function buildObjectSidebarMenu(context: SidebarMenuFactoryContext): boolean {
     const destructiveActions: ContextMenuItem[] = [];
     items.push(copyNameMenuItem());
     items.push({ label: t("contextMenu.newQuery"), action: newQuery, icon: TerminalSquare });
+    if (node.type === "table" && supportsAiAssistantContext(currentDatabaseType())) {
+      items.push({ label: t("contextMenu.addToAi"), action: () => emit("add-to-ai", node), icon: Sparkles });
+    }
     items.push({ label: "", separator: true });
     items.push({ label: t("contextMenu.viewData"), action: openDataImmediately, icon: TableProperties });
     items.push({

apps/desktop/src/i18n/locales/en.ts

@@ -2698,6 +2698,7 @@ export default {
     duplicateConnection: "Duplicate Connection",
     duplicateSelectedConnections: "Duplicate Selected {count} Connections",
     newQuery: "New Query",
+    addToAi: "Add to AI",
     instanceInfo: "Instance Info",
     newSql: "New SQL",
     generateSql: "Generate SQL",

apps/desktop/src/i18n/locales/es.ts

@@ -2631,6 +2631,7 @@ export default withEnglishFallback({
     duplicateSelectedConnections: "Duplicar {count} conexiones seleccionadas",
     instanceInfo: "Información de la instancia",
     newQuery: "Nueva consulta",
+    addToAi: "Agregar a la IA",
     newSql: "Nuevo SQL",
     generateSql: "Generar SQL",
     newInsert: "Nueva Inserción",

apps/desktop/src/i18n/locales/it.ts

@@ -2628,6 +2628,7 @@ export default withEnglishFallback({
     duplicateConnection: "Duplica Connessione",
     duplicateSelectedConnections: "Duplica {count} connessioni selezionate",
     newQuery: "Nuova Query",
+    addToAi: "Aggiungi all'IA",
     instanceInfo: "Informazioni Istanza",
     newSql: "Nuovo SQL",
     generateSql: "Genera SQL",

apps/desktop/src/i18n/locales/ja.ts

@@ -2654,6 +2654,7 @@ export default withEnglishFallback({
     duplicateConnection: "接続を複製",
     duplicateSelectedConnections: "選択した{count}件の接続を複製",
     newQuery: "新しいクエリ",
+    addToAi: "AI に追加",
     newSql: "新しいSQL",
     generateSql: "SQLを生成",
     newInsert: "新しいINSERT",

apps/desktop/src/i18n/locales/ko.ts

@@ -2542,6 +2542,7 @@ export default withEnglishFallback({
     duplicateConnection: "연결 복제",
     duplicateSelectedConnections: "선택한 연결 {count}개 복제",
     newQuery: "새 쿼리",
+    addToAi: "AI에 추가",
     instanceInfo: "인스턴스 정보",
     newSql: "새 SQL",
     generateSql: "SQL 생성",

apps/desktop/src/i18n/locales/pt-BR.ts

@@ -2631,6 +2631,7 @@ export default withEnglishFallback({
     duplicateSelectedConnections: "Duplicar {count} conexões selecionadas",
     instanceInfo: "Informações da Instância",
     newQuery: "Nova Consulta",
+    addToAi: "Adicionar à IA",
     newSql: "Novo SQL",
     generateSql: "Gerar SQL",
     newInsert: "Novo Insert",

apps/desktop/src/i18n/locales/zh-CN.ts

@@ -2622,6 +2622,7 @@ export default withEnglishFallback({
     damengJobAdmin: "达梦代理作业",
     openDamengJobAdmin: "打开达梦代理作业",
     newQuery: "新建查询",
+    addToAi: "添加到 AI",
     instanceInfo: "实例信息",
     newInsert: "新建新增",
     newSql: "新建 SQL",

apps/desktop/src/i18n/locales/zh-TW.ts

@@ -2629,6 +2629,7 @@ export default withEnglishFallback({
     duplicateConnection: "複製連線",
     duplicateSelectedConnections: "複製選取的 {count} 個連線",
     newQuery: "建立查詢",
+    addToAi: "加入 AI",
     instanceInfo: "實例資訊",
     newSql: "新增 SQL",
     generateSql: "產生 SQL",

apps/desktop/src/lib/__tests__/ai/addToAiDatabaseContext.spec.ts

@@ -0,0 +1,15 @@
+import { readFileSync } from "node:fs";
+import { describe, expect, it } from "vitest";
+
+const appSource = readFileSync(new URL("../../../App.vue", import.meta.url), "utf8");
+
+describe("add to AI database context", () => {
+  it("accepts database tree nodes with an empty database identifier", () => {
+    const addToAiStart = appSource.indexOf("async function addToAi");
+    const addToAiEnd = appSource.indexOf("function openAiPanel", addToAiStart);
+    const addToAiSource = appSource.slice(addToAiStart, addToAiEnd);
+
+    expect(addToAiSource).toContain("else if (hasTreeNodeDatabaseContext(node))");
+    expect(addToAiSource).not.toContain("else if (node.database)");
+  });
+});

apps/desktop/src/lib/database/databaseFeatureSupport.ts

@@ -92,6 +92,15 @@ export function supportsQueryExecution(dbType?: DatabaseType): boolean {
   return supportsDatabaseFeature(dbType, "queryExecution");
 }
 
+/**
+ * The AI assistant currently builds its context from database/table metadata.
+ * Connection-only query targets (for example etcd and ZooKeeper) do not expose
+ * that hierarchy, so they must not be offered by sidebar "Add to AI" actions.
+ */
+export function supportsAiAssistantContext(dbType?: DatabaseType): boolean {
+  return supportsQueryExecution(dbType) && !usesConnectionOnlyQueryTarget(dbType);
+}
+
 export function supportsConnectionScopedQueryExecution(dbType?: DatabaseType): boolean {
   return supportsRegisteredConnectionScopedQueryExecution(dbType);
 }