From 193d875b8abfde2764df5d5e606a43cb96d16192 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BA=8C=E4=B8=AB=E8=AE=B2=E6=A2=B5?= Date: Mon, 14 Sep 2026 13:50:17 +0800 Subject: [PATCH] feat(ai): cross-database query with multi database selection --- .../src/components/editor/AiAssistant.vue | 109 +++++++++++------- apps/desktop/src/lib/__tests__/ai/ai.spec.ts | 30 ++++- apps/desktop/src/lib/ai/ai.ts | 16 ++- apps/desktop/src/lib/backend/http.ts | 2 + apps/desktop/src/lib/backend/tauri.ts | 2 + crates/dbx-core/src/agent_loop.rs | 2 + crates/dbx-core/src/agent_tools.rs | 66 +++++++++-- crates/dbx-core/src/ai_claude_code_cli.rs | 1 + crates/dbx-core/src/ai_cli_agent.rs | 49 +++++++- crates/dbx-core/src/ai_codebuddy_cli.rs | 3 + crates/dbx-core/src/ai_codex_cli.rs | 1 + crates/dbx-core/src/ai_cursor_cli.rs | 2 + crates/dbx-core/src/ai_grok_cli.rs | 1 + crates/dbx-core/src/ai_opencode_cli.rs | 2 + crates/dbx-core/src/ai_pi_agent_cli.rs | 1 + crates/dbx-core/src/ai_qoder_cli.rs | 1 + .../tests/live_mongodb_agent_tools.rs | 2 + crates/dbx-mcp/src/lib.rs | 2 + crates/dbx-web/src/routes/ai.rs | 3 + src-tauri/src/commands/ai.rs | 2 + 20 files changed, 240 insertions(+), 57 deletions(-) diff --git a/apps/desktop/src/components/editor/AiAssistant.vue b/apps/desktop/src/components/editor/AiAssistant.vue index 6d4cf9dcca..f4fc18dbc7 100644 --- a/apps/desktop/src/components/editor/AiAssistant.vue +++ b/apps/desktop/src/components/editor/AiAssistant.vue @@ -50,7 +50,6 @@ import { } from "@lucide/vue"; import { Button } from "@/components/ui/button"; import { DropdownMenu, DropdownMenuContent, DropdownMenuItem, DropdownMenuTrigger } from "@/components/ui/dropdown-menu"; -import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Popover, PopoverAnchor, PopoverContent, PopoverTrigger } from "@/components/ui/popover"; import { ScrollArea } from "@/components/ui/scroll-area"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; @@ -156,7 +155,7 @@ import type { AiConfigItem, AiEffortCapability, AiEffortOption, AiEffortSelectio import type { ConnectionConfig, QueryTab, SavedSqlFile, TableInfo } from "@/types/database"; import { fetchNamespaceOptionsForConnection, useDatabaseOptions } from "@/composables/useDatabaseOptions"; import { useSchemaOptions } from "@/composables/useSchemaOptions"; -import { decodeSelectableDatabaseValue, encodeSelectableDatabaseValue, formatDatabaseLabel, resolveDefaultDatabase } from "@/lib/database/defaultDatabase"; +import { encodeSelectableDatabaseValue, formatDatabaseLabel, resolveDefaultDatabase } from "@/lib/database/defaultDatabase"; import { normalizeSqliteNamespace } from "@/lib/database/sqliteNamespace"; import { isQueryExecutionErrorResult } from "@/lib/query/queryResultError"; import { isSchemaAware, isSingleDatabase } from "@/lib/database/databaseCapabilities"; @@ -1405,19 +1404,43 @@ const dbSelectOptions = computed(() => { })); }); -const selectedNamespace = computed(() => (props.connection && props.tab ? resolveAiNamespaceSelection(props.tab, props.connection).value : "")); - -const selectedDatabaseSelectValue = computed(() => (props.connection ? encodeSelectableDatabaseValue(props.connection.db_type, selectedNamespace.value) : "")); +// AI can inspect more than the tab's active database. Keep this selection local +// to the composer so changing the visible query tab does not rewrite SQL state. +const selectedDatabases = ref([]); +const selectedDatabaseValues = computed(() => new Set(selectedDatabases.value)); const selectedDatabaseLabel = computed(() => { if (!props.connection) return t("editor.selectDatabase"); - if (!props.tab) return t("editor.selectDatabase"); - return formatDatabaseLabel(props.connection, selectedNamespace.value, { - defaultDatabase: t("editor.defaultDatabase"), - noDatabase: t("editor.noDatabase"), - }); + const labels = dbSelectOptions.value.filter((option) => selectedDatabaseValues.value.has(option.database)).map((option) => option.label); + return labels.length ? labels.join(", ") : t("editor.selectDatabase"); }); +function syncSelectedDatabases() { + const active = selectedNamespace.value; + const available = dbSelectOptions.value.map((option) => option.database); + const retained = selectedDatabases.value.filter((database) => available.includes(database)); + selectedDatabases.value = retained.length ? retained : active ? [active] : []; +} + +function toggleDatabase(database: string) { + if (selectedDatabaseValues.value.has(database)) { + if (selectedDatabases.value.length === 1) return; + selectedDatabases.value = selectedDatabases.value.filter((item) => item !== database); + } else { + selectedDatabases.value = [...selectedDatabases.value, database]; + } +} + +const selectedNamespace = computed(() => (props.connection && props.tab ? resolveAiNamespaceSelection(props.tab, props.connection).value : "")); + +watch( + () => `${props.connection?.id ?? ""}:${props.tab?.id ?? ""}`, + () => { + selectedDatabases.value = selectedNamespace.value ? [selectedNamespace.value] : []; + }, +); +watch([dbSelectOptions, selectedNamespace], syncSelectedDatabases, { immediate: true }); + const showAiSchemaSelector = computed(() => { const connection = props.connection; return !!connection && connection.db_type !== "dameng" && aiSchemaSelectionSupported(connection); @@ -1478,20 +1501,6 @@ async function changeConnection(connectionId: string) { } } -function changeNamespace(value: string) { - const tab = props.tab; - 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 { - queryStore.updateDatabase(tab.id, namespace); - } -} - function changeSchema(schema: string) { const tab = props.tab; if (!tab || tab.schema === schema) return; @@ -2793,6 +2802,9 @@ async function send() { clearPendingWriteGrant(); return; } + // Capture the selection before context loading or queued run scheduling can + // yield to another conversation. Dameng's top-level selector is a schema. + const runDatabases = resolveAiNamespaceSelection(tab, connection).kind === "database" ? [...selectedDatabases.value] : []; const activeConfig = activeFullConfig.value; if (!activeConfig) { clearPendingWriteGrant(); @@ -3086,11 +3098,21 @@ async function send() { // paying for buildAiContext() too; it can do real backend/schema work that // would be entirely wasted on an already-abandoned request. if (!generationCanContinue()) return; - const context = await buildAiContext(tab, connection, { - mentionedTables, - sqlFiles, - csvFiles: csvAttachments, - }); + const requestDatabase = runDatabases[0] ?? tab.database; + const context = await buildAiContext( + { + ...tab, + database: requestDatabase, + schema: runDatabases.length > 1 || requestDatabase !== tab.database ? undefined : tab.schema, + }, + connection, + { + mentionedTables, + sqlFiles, + csvFiles: csvAttachments, + }, + ); + context.selectedDatabases = runDatabases; // Superseded while awaiting buildAiContext() above — must bail before ever // calling runAgentStream(), not just before writing its results. Without // this recheck, a clear/switch/unmount that fires during context @@ -4970,27 +4992,26 @@ async function openExternalUrl(url: string) { />