feat(ai): cross-database query with multi database selection

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<string[]>([]);
 
+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) {
               />
               <template v-if="connection">
                 <Database class="h-3 w-3 shrink-0 text-foreground/40" />
-                <Select
-                  :model-value="selectedDatabaseSelectValue"
-                  @update:model-value="
-                    (v) => {
-                      if (typeof v === 'string') changeNamespace(v);
-                    }
-                  "
+                <Popover
                   @update:open="
                     (open: boolean) => {
                       if (open) loadDatabases();
                     }
                   "
                 >
-                  <SelectTrigger :class="['h-5 w-auto border-0 rounded-md bg-transparent dark:bg-transparent p-0 px-1 text-xs text-foreground/80 shadow-none focus:ring-0 focus-visible:ring-0 [&_svg]:size-3', showAiSchemaSelector && 'min-w-0 max-w-56 flex-1']">
-                    <SelectValue :placeholder="t('editor.selectDatabase')">{{ selectedDatabaseLabel }}</SelectValue>
-                  </SelectTrigger>
-                  <SelectContent>
-                    <SelectItem v-for="option in dbSelectOptions" :key="option.value" :value="option.value">{{ option.label }}</SelectItem>
-                    <SelectItem v-if="!dbSelectOptions.length && connection && tab" :value="selectedDatabaseSelectValue">{{ selectedDatabaseLabel }}</SelectItem>
-                  </SelectContent>
-                </Select>
+                  <PopoverTrigger as-child>
+                    <Button variant="ghost" :class="['h-5 max-w-64 justify-start border-0 p-0 px-1 text-xs font-normal text-foreground/80 shadow-none', showAiSchemaSelector && 'min-w-0 flex-1']">
+                      <span class="truncate">{{ selectedDatabaseLabel }}</span>
+                    </Button>
+                  </PopoverTrigger>
+                  <PopoverContent align="start" class="w-64 p-1">
+                    <button v-for="option in dbSelectOptions" :key="option.value" type="button" class="flex w-full items-center gap-2 rounded-sm px-2 py-1.5 text-left text-sm hover:bg-muted" @click="toggleDatabase(option.database)">
+                      <Check :class="['h-4 w-4', selectedDatabaseValues.has(option.database) ? 'opacity-100' : 'opacity-0']" />
+                      <span class="truncate">{{ option.label }}</span>
+                    </button>
+                    <div v-if="!dbSelectOptions.length" class="px-2 py-1.5 text-sm text-muted-foreground">{{ t("editor.selectDatabase") }}</div>
+                  </PopoverContent>
+                </Popover>
                 <template v-if="showAiSchemaSelector">
                   <Layers class="h-3 w-3 shrink-0 text-foreground/40" />
                   <SearchableSelect

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

@@ -1,5 +1,6 @@
-import { beforeAll, describe, expect, it } from "vitest";
-import { buildAgentRequest, buildSystemPrompt, buildUserPrompt, type AiContext } from "@/lib/ai/ai";
+import { beforeAll, describe, expect, it, vi } from "vitest";
+import { buildAgentRequest, buildSystemPrompt, buildUserPrompt, runAgentStream, type AiContext } from "@/lib/ai/ai";
+import * as api from "@/lib/backend/api";
 import { setLocale } from "@/i18n";
 
 function context(overrides: Partial<AiContext> = {}): AiContext {
@@ -18,6 +19,31 @@ function context(overrides: Partial<AiContext> = {}): AiContext {
 }
 
 describe("AI SQL dialect prompt", () => {
+  it("sends database selections to the backend as well as the model", async () => {
+    const stream = vi.spyOn(api, "aiAgentStream").mockResolvedValue("done");
+    try {
+      await runAgentStream(
+        {
+          config: { provider: "openai", apiKey: "test", apiUrl: "https://example.invalid", model: "model" },
+          action: "general",
+          mode: "agent",
+          instruction: "Join users and orders",
+          context: context({ databaseType: "mysql", database: "db_a", selectedDatabases: ["db_a", "db_b"] }),
+        },
+        [],
+        () => {},
+        "multi-db-run",
+      );
+      const args = stream.mock.calls[0];
+      expect(args[3]).toBe("db_a");
+      expect(args[14]).toEqual(["db_a", "db_b"]);
+      expect(args[1].systemPrompt).toContain('Selected databases: ["db_a","db_b"]');
+      expect(args[8]).toBe(false);
+    } finally {
+      stream.mockRestore();
+    }
+  });
+
   // buildSystemPrompt picks zh/en copy via currentLocale(); pin to en so the
   // English-string assertions are deterministic regardless of the host OS locale.
   beforeAll(async () => {

apps/desktop/src/lib/ai/ai.ts

@@ -100,6 +100,8 @@ export interface AiContext {
   connectionName: string;
   databaseType: DatabaseType;
   database: string;
+  /** Databases selected for this run; omitted by older callers. */
+  selectedDatabases?: string[];
   /** Schema selected for metadata loading and agent tool execution. */
   schema?: string;
   currentSql: string;
@@ -221,7 +223,8 @@ export async function runAgentStream(input: AiRequestInput, history: api.AiMessa
   const { messages, systemPrompt, taskContract, maxTokens } = buildAgentRequest(input, history, custom);
   const sid = sessionId || uuid();
 
-  return api.aiAgentStream(
+  const selectedDatabases = input.context.selectedDatabases;
+  const args = [
     sid,
     {
       config: input.config,
@@ -241,7 +244,11 @@ export async function runAgentStream(input: AiRequestInput, history: api.AiMessa
     input.confirmedConnectionId,
     input.confirmedDatabase,
     input.confirmedSchema,
-  );
+  ] as const;
+  if (selectedDatabases?.length) {
+    return api.aiAgentStream(...args, undefined, selectedDatabases);
+  }
+  return api.aiAgentStream(...args);
 }
 
 export function buildUserPrompt(action: AiAction, context: AiContext, instruction: string, isZh: boolean): string {
@@ -319,6 +326,11 @@ export function buildSystemPrompt(action: AiAction, context: AiContext, mode: Ai
     `Database type: ${context.databaseType}`,
     `Connection: ${context.connectionName}`,
     `Database: ${context.database}`,
+    context.selectedDatabases?.length
+      ? isZh
+        ? `已选择数据库：${JSON.stringify(context.selectedDatabases)}。请在元数据工具中通过 database 参数分别检查这些数据库。跨库 JOIN 需要数据库引擎支持：MySQL 使用 database.table，SQL Server 使用 database.schema.table；PostgreSQL 不能直接连接不同数据库，除非已配置联邦查询。MCP 授权仍然生效，选择数据库不会绕过授权限制。`
+        : `Selected databases: ${JSON.stringify(context.selectedDatabases)}. Use the database parameter on metadata tools to inspect each selected database. Cross-database joins require engine support: MySQL uses database.table, SQL Server uses database.schema.table; PostgreSQL cannot directly join separate databases without an existing federation setup. MCP authorization still applies; selecting databases does not override it.`
+      : "",
     context.schema ? `Selected schema: ${context.schema}` : "",
     schemaCoverageLine(context, isZh),
     "",

apps/desktop/src/lib/backend/http.ts

@@ -1873,6 +1873,7 @@ export async function aiAgentStream(
   confirmedDatabase?: string,
   confirmedSchema?: string,
   signal?: AbortSignal,
+  selectedDatabases?: string[],
 ): Promise<string> {
   const res = await fetch(apiUrl("/api/ai/agent-stream"), {
     method: "POST",
@@ -1890,6 +1891,7 @@ export async function aiAgentStream(
       confirmedConnectionId,
       confirmedDatabase,
       confirmedSchema,
+      selectedDatabases,
     }),
     signal,
   });

apps/desktop/src/lib/backend/tauri.ts

@@ -587,6 +587,7 @@ export async function aiAgentStream(
   confirmedDatabase?: string,
   confirmedSchema?: string,
   _signal?: AbortSignal,
+  selectedDatabases?: string[],
 ): Promise<string> {
   const unlisten: UnlistenFn = await listen<TauriAgentEvent>("ai-agent-event", (event) => {
     const payload = event.payload;
@@ -610,6 +611,7 @@ export async function aiAgentStream(
       confirmedConnectionId,
       confirmedDatabase,
       confirmedSchema,
+      selectedDatabases,
     });
   } catch (e) {
     unlisten();