diff --git a/apps/desktop/src/components/editor/AiAssistant.vue b/apps/desktop/src/components/editor/AiAssistant.vue index 342ab22ed3..b861296e1d 100644 --- a/apps/desktop/src/components/editor/AiAssistant.vue +++ b/apps/desktop/src/components/editor/AiAssistant.vue @@ -73,6 +73,7 @@ import { useNavigationTargets } from "@/composables/useNavigationTargets"; import { buildAiContext, resolveAiDatabaseTarget, + resolveAiMentionDatabase, resolveAiNamespaceSelection, resolveDefaultAiSchema, aiDatabaseTypeForConnection, @@ -2072,10 +2073,15 @@ function normalizeMentionQuery(query: string): { schemaPrefix: string; tableFilt }; } -async function loadMentionCandidates(query: string) { - if (!props.connection || !props.tab?.connectionId || !props.tab.database) return; +function mentionTargetDatabase(): string { + return props.tab && props.connection ? resolveAiMentionDatabase(props.tab, props.connection, selectedDatabases.value) : ""; +} - const key = mentionCacheKey(props.tab.connectionId, props.tab.database, query); +async function loadMentionCandidates(query: string) { + const mentionDatabase = mentionTargetDatabase(); + if (!props.connection || !props.tab?.connectionId || !mentionDatabase) return; + + const key = mentionCacheKey(props.tab.connectionId, mentionDatabase, query); if (mentionCache.value[key]) { mentionCandidates.value = mentionCache.value[key]; return; @@ -2092,11 +2098,11 @@ async function loadMentionCandidates(query: string) { await connectionStore.ensureConnected(props.tab.connectionId); let tableCandidates: AiMentionCandidate[] = []; if (isSchemaAware(props.connection.db_type)) { - const schemas = mentionSchemaOrder(await listSchemas(props.tab.connectionId, props.tab.database)); + const schemas = mentionSchemaOrder(await listSchemas(props.tab.connectionId, mentionDatabase)); const filteredSchemas = schemaPrefix ? schemas.filter((schema) => schema.toLowerCase().includes(schemaPrefix.toLowerCase())) : schemas; const results = await Promise.all( filteredSchemas.slice(0, AI_TABLE_MENTION_SCHEMA_LIMIT).map(async (schema) => { - const tables = await listTables(props.tab!.connectionId, props.tab!.database, schema, tableFilter || undefined, AI_TABLE_MENTION_CANDIDATE_LIMIT); + const tables = await listTables(props.tab!.connectionId, mentionDatabase, schema, tableFilter || undefined, AI_TABLE_MENTION_CANDIDATE_LIMIT); return filterAiTableMentionCandidates( tables.map((table) => mentionCandidateFromTable(table, schema)), tableFilter, @@ -2106,7 +2112,7 @@ async function loadMentionCandidates(query: string) { ); tableCandidates = filterAiTableMentionCandidates(results.flat(), "", AI_TABLE_MENTION_CANDIDATE_LIMIT); } else { - const database = props.connection.db_type === "sqlite" ? normalizeSqliteNamespace(props.tab.database || props.connection.database, props.connection) : props.tab.database; + const database = props.connection.db_type === "sqlite" ? normalizeSqliteNamespace(mentionDatabase || props.connection.database, props.connection) : mentionDatabase; const schema = database || props.connection.database || "main"; const tables = await listTables(props.tab.connectionId, database, schema, tableFilter || undefined, AI_TABLE_MENTION_CANDIDATE_LIMIT); tableCandidates = filterAiTableMentionCandidates( @@ -2340,7 +2346,7 @@ function imageAttachmentSupportErrorMessage(error: "provider" | "format"): strin function selectedMessageMentions(tableMentions: AiTableMention[], sqlFileMentions: AiSqlFileMention[], csvAttachments: AiCsvFileContext[] = [], imageAttachments: AiImageAttachment[] = []): AiMessageMention[] { const connectionId = props.tab?.connectionId || props.connection?.id || ""; - const database = props.tab?.database || props.connection?.database || ""; + const database = mentionTargetDatabase() || props.connection?.database || ""; return [ ...tableMentions.map((mention) => ({ kind: "table" as const, @@ -2441,7 +2447,7 @@ function refreshMentionState() { commandOpen.value = false; const mention = activeMentionAtCursor(); - if (!mention || !props.connection || !props.tab?.database) { + if (!mention || !props.connection || !mentionTargetDatabase()) { mentionOpen.value = false; return; } @@ -2453,6 +2459,10 @@ function refreshMentionState() { }, 120); } +watch(selectedDatabases, () => { + if (mentionOpen.value) refreshMentionState(); +}); + function onPromptKeyup(event: KeyboardEvent) { if (["ArrowDown", "ArrowUp", "Enter", "Tab", "Escape"].includes(event.key)) return; refreshMentionState(); diff --git a/apps/desktop/src/lib/__tests__/ai/aiContext.spec.ts b/apps/desktop/src/lib/__tests__/ai/aiContext.spec.ts index a1121b875a..1030e303a9 100644 --- a/apps/desktop/src/lib/__tests__/ai/aiContext.spec.ts +++ b/apps/desktop/src/lib/__tests__/ai/aiContext.spec.ts @@ -1,5 +1,5 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; -import { aiSchemaSelectionSupported, buildAiContext, resolveAiDatabaseTarget, resolveAiNamespaceSelection, resolveDefaultAiSchema, runAgentStream } from "@/lib/ai/ai"; +import { aiSchemaSelectionSupported, buildAiContext, resolveAiDatabaseTarget, resolveAiMentionDatabase, resolveAiNamespaceSelection, resolveDefaultAiSchema, runAgentStream } from "@/lib/ai/ai"; import type { AiConfig } from "@/types/ai"; import type { ConnectionConfig, QueryTab } from "@/types/database"; @@ -145,6 +145,14 @@ describe("Dameng AI context routing", () => { expect(resolveDefaultAiSchema(connection, ["REPORTING", "ARCHIVE"])).toBe("REPORTING"); }); + it("lists @ mention tables from the database picked in the composer", () => { + const connection = postgresConnection(); + expect(resolveAiMentionDatabase(queryTab("db_first"), connection, ["db_second"])).toBe("db_second"); + expect(resolveAiMentionDatabase(queryTab("db_first"), connection, ["db_second", "db_first"])).toBe("db_second"); + expect(resolveAiMentionDatabase(queryTab("db_first"), connection, [])).toBe("db_first"); + expect(resolveAiMentionDatabase(queryTab("db_first"), damengConnection(), ["ignored"])).toBe("db_first"); + }); + it("does not change non-Dameng namespace behavior", () => { expect(resolveAiDatabaseTarget(queryTab("analytics"), sqliteConnection())).toEqual({ database: "analytics" }); }); diff --git a/apps/desktop/src/lib/ai/ai.ts b/apps/desktop/src/lib/ai/ai.ts index 534d322979..b61ab33c88 100644 --- a/apps/desktop/src/lib/ai/ai.ts +++ b/apps/desktop/src/lib/ai/ai.ts @@ -820,6 +820,17 @@ export function resolveAiNamespaceSelection(tab: QueryTab, connection: Connectio return { kind: "database", value: tab.database || "" }; } +/** + * Database that `@` table mentions are listed from and resolved against. It has + * to match the request database (`selectedDatabases[0] ?? tab.database`), so a + * database picked in the composer wins over the query tab's own database. + */ +export function resolveAiMentionDatabase(tab: QueryTab, connection: ConnectionConfig, selectedDatabases: string[]): string { + const namespace = resolveAiNamespaceSelection(tab, connection); + if (namespace.kind !== "database") return tab.database || ""; + return selectedDatabases[0] ?? tab.database ?? ""; +} + export function resolveDefaultAiSchema(connection: ConnectionConfig, schemaOptions: string[]): string | undefined { if (connection.db_type !== "dameng") return undefined; const username = connection.username.trim().toLowerCase();