From 91f309f6c7134e0aa12a0d720f83ad3d5c0ee51e Mon Sep 17 00:00:00 2001 From: Abeautifulsnow Date: Wed, 2 Sep 2026 20:37:39 +0800 Subject: [PATCH] feat(ai): default and last-used templates per db type --- .../src/components/editor/AiAssistant.vue | 93 ++++++++++++-- .../editor/EditorSettingsDialog.vue | 54 ++++++++ .../AiAssistant.templateDefaults.spec.ts | 80 ++++++++++++ apps/desktop/src/i18n/locales/en.ts | 4 + apps/desktop/src/i18n/locales/es.ts | 4 + apps/desktop/src/i18n/locales/it.ts | 4 + apps/desktop/src/i18n/locales/ja.ts | 4 + apps/desktop/src/i18n/locales/ko.ts | 4 + apps/desktop/src/i18n/locales/pt-BR.ts | 4 + apps/desktop/src/i18n/locales/zh-CN.ts | 4 + apps/desktop/src/i18n/locales/zh-TW.ts | 4 + .../ai/promptTemplateDefaults.spec.ts | 88 +++++++++++++ apps/desktop/src/lib/ai/ai.ts | 2 +- .../src/lib/ai/promptTemplateDefaults.ts | 58 +++++++++ ...ttingsStore.promptTemplateDefaults.spec.ts | 119 ++++++++++++++++++ apps/desktop/src/stores/settingsStore.ts | 79 ++++++++++++ apps/desktop/src/types/ai.ts | 4 + crates/dbx-core/src/ai.rs | 12 +- crates/dbx-core/src/storage.rs | 36 ++++++ docs/content/docs/ai-assistant.mdx | 2 + 20 files changed, 647 insertions(+), 12 deletions(-) create mode 100644 apps/desktop/src/components/editor/__tests__/AiAssistant.templateDefaults.spec.ts create mode 100644 apps/desktop/src/lib/__tests__/ai/promptTemplateDefaults.spec.ts create mode 100644 apps/desktop/src/lib/ai/promptTemplateDefaults.ts create mode 100644 apps/desktop/src/stores/__tests__/settingsStore.promptTemplateDefaults.spec.ts diff --git a/apps/desktop/src/components/editor/AiAssistant.vue b/apps/desktop/src/components/editor/AiAssistant.vue index fb02efc757..92e692bd87 100644 --- a/apps/desktop/src/components/editor/AiAssistant.vue +++ b/apps/desktop/src/components/editor/AiAssistant.vue @@ -38,6 +38,7 @@ import { Table2, Play, Square, + Star, Trash2, Terminal, Wand2, @@ -73,6 +74,7 @@ import { resolveAiDatabaseTarget, resolveAiNamespaceSelection, resolveDefaultAiSchema, + aiDatabaseTypeForConnection, aiSchemaSelectionSupported, runAgentStream, isVectorDbType, @@ -135,6 +137,8 @@ import { orderAiConfigsForDisplay } from "@/lib/ai/aiConfigOrdering"; import { effortSelectionEquals, runtimeEffortFromPreference } from "@/lib/ai/aiEffortPreference"; import { useAiModelCatalog } from "@/composables/useAiModelCatalog"; import { ACTIVE_TEMPLATES_TOTAL_MAX, promptTemplateCharacterCount } from "@/types/promptTemplate"; +import { capTemplateIdsToCharLimit, resolveAutoTemplateIds, resolveDefaultTemplateIds } from "@/lib/ai/promptTemplateDefaults"; +import { databaseManifestEntry } from "@/lib/database/databaseDriverManifest"; import type { AgentEvent } from "@/lib/backend/tauri"; import { buildAiAgentPlan } from "@/lib/ai/aiAgentPlan"; @@ -438,19 +442,70 @@ watch( // Retry store load on selector open if prior init failed (e.g. backend not yet ready at mount) watch(showTemplateSelector, (open) => { - if (open) void promptTemplateStore.ensureLoaded(); + // Retry load on selector open if prior init failed, then apply pending + // per-db_type defaults that were skipped when the initial load failed. + if (open) void promptTemplateStore.ensureLoaded().then(() => void maybeApplyAutoTemplates()); }); -// Reset template selection when the user switches to a different connection, -// database, or schema — a new namespace context warrants a fresh selection. +// Auto-apply per-db_type prompt template defaults once both the AI selection +// (which carries the defaults) and the template list have loaded. Only this +// one-time resolution and the namespace watcher below write the selection +// implicitly; manual edits in between are never overwritten. +// Template defaults/last-used are keyed by the connection's effective AI +// database type (gbase → mysql/informix, doris-over-mysql protocol, jdbc → +// inferred dialect) so resolution matches the dialect the AI pipeline and +// prompt selection actually use — the same axis aiDatabaseTypeForConnection +// established for schema selection. +const templateDbType = computed(() => (props.connection ? aiDatabaseTypeForConnection(props.connection) : undefined)); +const aiTemplateNamespaceKey = computed(() => `${props.connection?.id ?? ""}::${props.tab?.database ?? ""}::${props.tab?.schema ?? ""}`); +let autoTemplatesInitialized = false; +function applyResolvedTemplateIds(ids: string[]) { + activeTemplateIds.value = capTemplateIdsToCharLimit(ids, promptTemplateStore.templates, ACTIVE_TEMPLATES_TOTAL_MAX); +} +async function maybeApplyAutoTemplates() { + if (autoTemplatesInitialized || !settings.isAiConfigLoaded) return; + const namespaceAtStart = aiTemplateNamespaceKey.value; + if (!(await promptTemplateStore.ensureLoaded())) return; + autoTemplatesInitialized = true; + // A namespace switch during the await already ran the defaults-only + // resolution below; applying defaults-or-last-used now would violate the + // switch contract and could overwrite a manual selection made meanwhile. + if (namespaceAtStart !== aiTemplateNamespaceKey.value) return; + // Panel-open resolution: explicit defaults, else the db_type's last-used. + applyResolvedTemplateIds( + resolveAutoTemplateIds({ + dbType: templateDbType.value, + defaultTemplatesByDbType: settings.aiDefaultTemplatesByDbType, + lastUsedTemplatesByDbType: settings.aiLastUsedTemplatesByDbType, + }), + ); +} watch( - // Return a stable primitive key: a fresh array literal is never Object.is-equal to the - // previous one, so a getter returning `[id, database]` fires on every dependency - // invalidation (e.g. the 30s backup scheduler replacing connection objects) even when the - // id/database values are unchanged — spuriously clearing the selection mid agent-run. - () => `${props.connection?.id ?? ""}::${props.tab?.database ?? ""}::${props.tab?.schema ?? ""}`, + () => settings.isAiConfigLoaded, + () => void maybeApplyAutoTemplates(), + { immediate: true }, +); + +// Reset template selection when the user switches to a different connection, +// database, or schema — a new namespace context warrants a fresh selection: +// the new db_type's default templates when the user configured any (explicit +// opt-in), otherwise empty, preserving the pre-defaults clear-on-switch +// contract. Last-used templates are intentionally not restored here: that +// would silently re-select templates on every namespace switch. +watch( + // The key is a stable primitive string (see aiTemplateNamespaceKey): a fresh + // array literal is never Object.is-equal to the previous one, so a getter + // returning `[id, database]` fires on every dependency invalidation (e.g. + // the 30s backup scheduler replacing connection objects) even when the + // id/database values are unchanged — spuriously clearing the selection mid + // agent-run. + aiTemplateNamespaceKey, () => { - activeTemplateIds.value = []; + if (!settings.isAiConfigLoaded || !promptTemplateStore.isLoaded) { + activeTemplateIds.value = []; + return; + } + applyResolvedTemplateIds(resolveDefaultTemplateIds(templateDbType.value, settings.aiDefaultTemplatesByDbType)); }, ); @@ -489,6 +544,14 @@ const templateSelectorTriggerLabel = computed(() => { } return templateSelectorLabel.value; }); +const currentDbType = computed(() => templateDbType.value); +function isDefaultTemplateForCurrentDb(id: string): boolean { + const dbType = currentDbType.value; + return !!dbType && (settings.aiDefaultTemplatesByDbType[dbType]?.includes(id) ?? false); +} +function currentDbTypeLabel(): string { + return currentDbType.value ? (databaseManifestEntry(currentDbType.value)?.label ?? currentDbType.value) : ""; +} const promptTextareaRef = ref(null); const csvFileInputRef = ref(null); const shouldAutoScroll = ref(true); @@ -2823,6 +2886,13 @@ async function send() { globalInstructions: promptTemplateStore.globalInstructions, activeTemplates: [...activeTemplates.value], }; + // Remember what was actually sent for this db_type so panels opened later can + // restore it when no explicit per-db_type defaults are configured. This runs + // for empty selections too: the store clears the remembered entry so a + // deselected-everything send is not resurrected on the next panel open. + if (templateDbType.value) { + settings.recordLastUsedTemplates(templateDbType.value, [...activeTemplateIds.value]); + } const selectedTableMentions = auto ? [] : [...selectedMentions.value]; const selectedSqlFiles = auto ? [] : [...selectedSqlFileMentions.value]; @@ -4846,7 +4916,10 @@ async function openExternalUrl(url: string) {
-
{{ tpl.name }}
+
+ {{ tpl.name }} + +
{{ tpl.content.slice(0, 60) }}
diff --git a/apps/desktop/src/components/editor/EditorSettingsDialog.vue b/apps/desktop/src/components/editor/EditorSettingsDialog.vue index bcaeec06ac..8224f83f62 100644 --- a/apps/desktop/src/components/editor/EditorSettingsDialog.vue +++ b/apps/desktop/src/components/editor/EditorSettingsDialog.vue @@ -32,6 +32,7 @@ import { Search, Settings, Sun, + Star, SunMoon, Terminal, Trash2, @@ -40,6 +41,7 @@ import { } from "@lucide/vue"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; +import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover"; import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; import { Input } from "@/components/ui/input"; import PasswordInput from "@/components/ui/PasswordInput.vue"; @@ -207,6 +209,7 @@ import { MAX_QUERY_RESULT_MAX_ROWS } from "@/lib/dataGrid/queryResultRowLimit"; import type { PromptTemplate } from "@/types/promptTemplate"; import { GLOBAL_INSTRUCTIONS_MAX, PROMPT_TEMPLATE_CONTENT_MAX, PROMPT_TEMPLATE_NAME_MAX, promptTemplateCharacterCount } from "@/types/promptTemplate"; import { METADATA_CACHE_HARD_MAX_MEMORY_MB, METADATA_CACHE_MIN_MEMORY_MB, normalizeMetadataCacheMemoryMb } from "@/lib/metadata/metadataRuntimeCache"; +import { databaseManifestEntry, manifestDatabaseTypes } from "@/lib/database/databaseDriverManifest"; const { t } = useI18n(); const { toast } = useToast(); @@ -3353,6 +3356,8 @@ async function saveTemplateForm() { async function confirmDeleteTemplate(tpl: PromptTemplate) { try { await promptTemplateStore.remove(tpl.id); + // Keep the per-db_type default/last-used records free of the deleted id. + settingsStore.removeTemplateFromDefaultAndLastUsed(tpl.id); toast(t("ai.promptTemplateDeleted")); } catch (e: any) { toast(e?.message || String(e), 5000); @@ -3361,6 +3366,29 @@ async function confirmDeleteTemplate(tpl: PromptTemplate) { } } +// Per-db_type default templates (issue #7649): a template can be marked as the +// auto-applied default for any database type; the AI panel resolves these on +// mount/namespace switch. Stored in the AI chat selection, not on the template row. +const templateDefaultsOpenId = ref(""); +const dbTypeOptions = manifestDatabaseTypes(); +function defaultDbTypesForTemplate(templateId: string): string[] { + return Object.entries(settingsStore.aiDefaultTemplatesByDbType) + .filter(([, ids]) => ids.includes(templateId)) + .map(([dbType]) => dbType) + .sort(); +} +function templateHasDefault(dbType: string, templateId: string): boolean { + return settingsStore.aiDefaultTemplatesByDbType[dbType]?.includes(templateId) ?? false; +} +function toggleTemplateDefault(dbType: string, templateId: string) { + const current = settingsStore.aiDefaultTemplatesByDbType[dbType] ?? []; + const next = current.includes(templateId) ? current.filter((id) => id !== templateId) : [...current, templateId]; + settingsStore.setDefaultTemplatesForDbType(dbType, next); +} +function dbTypeLabel(dbType: string): string { + return databaseManifestEntry(dbType as DatabaseType)?.label ?? dbType; +} + async function saveGlobalInstructions() { if (promptTemplateCharacterCount(editGlobalInstructions.value) > GLOBAL_INSTRUCTIONS_MAX) return; globalInstructionsSaving.value = true; @@ -7324,8 +7352,34 @@ onUnmounted(() => { {{ tpl.name }}
{{ tpl.content.slice(0, 100) }}{{ tpl.content.length > 100 ? "..." : "" }}
+
+ + {{ dbTypeLabel(dbType) }} + +
+ + + + + +

{{ t("ai.templateDefaultsTitle") }}

+
+ +
+

+ {{ t("ai.templateDefaultsEmpty") }} +

+ +
diff --git a/apps/desktop/src/components/editor/__tests__/AiAssistant.templateDefaults.spec.ts b/apps/desktop/src/components/editor/__tests__/AiAssistant.templateDefaults.spec.ts new file mode 100644 index 0000000000..b6345aab03 --- /dev/null +++ b/apps/desktop/src/components/editor/__tests__/AiAssistant.templateDefaults.spec.ts @@ -0,0 +1,80 @@ +import { readFileSync } from "node:fs"; +import { describe, expect, it } from "vitest"; + +// Regression coverage for https://github.com/t8y2/dbx/issues/7649: per-db_type +// prompt template defaults + last-used fallback. The resolution rules live in +// lib/ai/promptTemplateDefaults.ts (unit-tested separately); this file pins +// that AiAssistant.vue wires them into the panel lifecycle: +// - auto-apply once when the AI selection + templates are loaded, +// - re-resolve on connection/database/schema switch (defaults only opt-in), +// - record what was actually sent as the db_type's last-used templates. +const source = readFileSync(new URL("../AiAssistant.vue", import.meta.url), "utf8"); + +describe("AI assistant applies per-db_type prompt template defaults", () => { + it("panel-open resolution uses defaults-else-last-used through the tested helper and enforces the char cap", () => { + const maybeBody = source.slice(source.indexOf("async function maybeApplyAutoTemplates()"), source.indexOf("// Reset template selection when the user switches")); + expect(maybeBody).toContain("resolveAutoTemplateIds({"); + expect(maybeBody).toContain("defaultTemplatesByDbType: settings.aiDefaultTemplatesByDbType"); + expect(maybeBody).toContain("lastUsedTemplatesByDbType: settings.aiLastUsedTemplatesByDbType"); + expect(maybeBody).toContain("applyResolvedTemplateIds("); + // The char cap lives in the shared apply helper both paths funnel through. + const applyBody = source.slice(source.indexOf("function applyResolvedTemplateIds("), source.indexOf("async function maybeApplyAutoTemplates()")); + expect(applyBody).toContain("capTemplateIdsToCharLimit(ids, promptTemplateStore.templates, ACTIVE_TEMPLATES_TOTAL_MAX)"); + }); + + it("auto-apply waits for the AI selection load, runs once, and skips failed template loads", () => { + const maybeBody = source.slice(source.indexOf("async function maybeApplyAutoTemplates()"), source.indexOf("// Reset template selection when the user switches")); + expect(maybeBody).toContain("if (autoTemplatesInitialized || !settings.isAiConfigLoaded) return;"); + expect(maybeBody).toContain("if (!(await promptTemplateStore.ensureLoaded())) return;"); + expect(maybeBody).toContain("autoTemplatesInitialized = true;"); + expect(maybeBody).toContain("applyResolvedTemplateIds("); + }); + + it("namespace switch re-resolves defaults only when both stores are loaded, otherwise clears", () => { + const watcherBody = source.slice(source.indexOf("aiTemplateNamespaceKey,"), source.indexOf("function toggleTemplateId(")); + expect(watcherBody).toContain("if (!settings.isAiConfigLoaded || !promptTemplateStore.isLoaded) {"); + expect(watcherBody).toContain("activeTemplateIds.value = [];"); + // The switch path must resolve through the defaults-only helper; the + // mount resolver (defaults-else-last-used) would resurrect cleared + // selections. Pin the exact call so an indirect last-used fallback + // cannot hide behind a shared helper again. + expect(watcherBody).toContain("resolveDefaultTemplateIds(templateDbType.value, settings.aiDefaultTemplatesByDbType)"); + expect(watcherBody).not.toContain("resolveAutoTemplateIds"); + expect(watcherBody).not.toContain("lastUsedTemplatesByDbType"); + }); + + it("keys defaults, last-used, and the badge by the effective AI database type", () => { + // Same axis aiDatabaseTypeForConnection established for schema selection: + // gbase → mysql, doris-over-mysql protocol → doris, jdbc → inferred + // dialect. Raw db_type keying would never match for those connections. + expect(source).toContain("const templateDbType = computed(() => (props.connection ? aiDatabaseTypeForConnection(props.connection) : undefined));"); + expect(source).toContain("aiDatabaseTypeForConnection"); + }); + + it("panel-open resolution bails when the namespace switched during the load await", () => { + const maybeBody = source.slice(source.indexOf("async function maybeApplyAutoTemplates"), source.indexOf("watch(", source.indexOf("async function maybeApplyAutoTemplates"))); + expect(maybeBody).toContain("const namespaceAtStart = aiTemplateNamespaceKey.value;"); + expect(maybeBody).toContain("if (namespaceAtStart !== aiTemplateNamespaceKey.value) return;"); + }); + + it("selector-open load retry also applies pending defaults", () => { + const watcherBody = source.slice(source.indexOf("watch(showTemplateSelector"), source.indexOf("// Auto-apply per-db_type")); + expect(watcherBody).toContain("void promptTemplateStore.ensureLoaded().then(() => void maybeApplyAutoTemplates());"); + }); + + it("send records the sent templates as the db_type's last-used selection", () => { + const sendIdx = source.indexOf("activeTemplates: [...activeTemplates.value],"); + const recordIdx = source.indexOf("settings.recordLastUsedTemplates(templateDbType.value, [...activeTemplateIds.value]);"); + expect(recordIdx).toBeGreaterThan(sendIdx); + // No length guard: an empty selection must also reach the store so the + // remembered entry is cleared instead of resurrecting old templates. + const guard = source.slice(sendIdx, recordIdx); + expect(guard).toContain("if (templateDbType.value) {"); + expect(guard).not.toContain("activeTemplates.value.length > 0"); + }); + + it("the template selector marks defaults for the current connection db_type", () => { + expect(source).toContain('v-if="isDefaultTemplateForCurrentDb(tpl.id)"'); + expect(source).toContain("settings.aiDefaultTemplatesByDbType[dbType]?.includes(id) ?? false"); + }); +}); diff --git a/apps/desktop/src/i18n/locales/en.ts b/apps/desktop/src/i18n/locales/en.ts index 8b9605204c..1b66ba8988 100644 --- a/apps/desktop/src/i18n/locales/en.ts +++ b/apps/desktop/src/i18n/locales/en.ts @@ -2945,6 +2945,10 @@ export default { templateSelectorLoading: "Loading...", templateSelectorEmpty: "No templates. Click Manage to add.", templateSelectorTooLong: "Selected templates exceed the total content limit ({max} characters). Deselect some to continue.", + templateSetDefault: "Set as default", + templateDefaultsTitle: "Default for database types", + templateDefaultsEmpty: "Not a default for any database type yet. Default templates are applied automatically when the AI panel opens on a matching connection.", + templateDefaultBadgeTitle: "Default template for {type}", exportMarkdown: "Export as Markdown", analysis: "AI Analysis", }, diff --git a/apps/desktop/src/i18n/locales/es.ts b/apps/desktop/src/i18n/locales/es.ts index f2414a4b98..fcf19413b9 100644 --- a/apps/desktop/src/i18n/locales/es.ts +++ b/apps/desktop/src/i18n/locales/es.ts @@ -2764,6 +2764,10 @@ export default withEnglishFallback({ templateSelectorLoading: "Cargando...", templateSelectorEmpty: "Sin plantillas. Haz clic en Gestionar para añadir.", templateSelectorTooLong: "Las plantillas seleccionadas superan el límite total de contenido ({max} caracteres). Deselecciona algunas para continuar.", + templateSetDefault: "Establecer como predeterminada", + templateDefaultsTitle: "Predeterminada por tipo de base de datos", + templateDefaultsEmpty: "Aún no es predeterminada para ningún tipo de base de datos. Las plantillas predeterminadas se aplican automáticamente al abrir el panel de IA en una conexión coincidente.", + templateDefaultBadgeTitle: "Plantilla predeterminada para {type}", exportMarkdown: "Exportar como Markdown", analysis: "AI Analysis", agentSteps: { diff --git a/apps/desktop/src/i18n/locales/it.ts b/apps/desktop/src/i18n/locales/it.ts index 6d9d43b0a3..ccf08fb66a 100644 --- a/apps/desktop/src/i18n/locales/it.ts +++ b/apps/desktop/src/i18n/locales/it.ts @@ -2658,6 +2658,10 @@ export default withEnglishFallback({ templateSelectorLoading: "Caricamento...", templateSelectorEmpty: "Nessun modello. Clicca Gestisci per aggiungerne.", templateSelectorTooLong: "I modelli selezionati superano il limite totale di contenuto ({max} caratteri). Deseleziona alcuni per continuare.", + templateSetDefault: "Imposta come predefinito", + templateDefaultsTitle: "Predefinito per tipo di database", + templateDefaultsEmpty: "Non ancora predefinito per alcun tipo di database. I modelli predefiniti vengono applicati automaticamente all'apertura del pannello AI su una connessione corrispondente.", + templateDefaultBadgeTitle: "Modello predefinito per {type}", exportMarkdown: "Esporta come Markdown", analysis: "AI Analysis", agentSteps: { diff --git a/apps/desktop/src/i18n/locales/ja.ts b/apps/desktop/src/i18n/locales/ja.ts index 9474c68d71..a7fb675e7f 100644 --- a/apps/desktop/src/i18n/locales/ja.ts +++ b/apps/desktop/src/i18n/locales/ja.ts @@ -2797,6 +2797,10 @@ export default withEnglishFallback({ templateSelectorLoading: "読み込み中...", templateSelectorEmpty: "テンプレートがありません。「管理」をクリックして追加してください。", templateSelectorTooLong: "選択されたテンプレートの合計が上限（{max}文字）を超えています。選択を解除してください。", + templateSetDefault: "デフォルトに設定", + templateDefaultsTitle: "データベース種別ごとのデフォルト", + templateDefaultsEmpty: "まだどのデータベース種別のデフォルトにも設定されていません。設定すると、該当するデータベースへの接続で AI パネルを開いたときに自動的に選択されます。", + templateDefaultBadgeTitle: "{type} のデフォルトテンプレート", exportMarkdown: "Markdown としてエクスポート", analysis: "AI 分析", agentSteps: { diff --git a/apps/desktop/src/i18n/locales/ko.ts b/apps/desktop/src/i18n/locales/ko.ts index 3e81e384d0..752a5b980c 100644 --- a/apps/desktop/src/i18n/locales/ko.ts +++ b/apps/desktop/src/i18n/locales/ko.ts @@ -2764,6 +2764,10 @@ export default withEnglishFallback({ templateSelectorLoading: "불러오는 중...", templateSelectorEmpty: "템플릿이 없습니다. 추가하려면 관리를 클릭하세요.", templateSelectorTooLong: "선택한 템플릿이 전체 내용 한도를 초과합니다 ({max}자). 계속하려면 일부를 선택 해제하세요.", + templateSetDefault: "기본값으로 설정", + templateDefaultsTitle: "데이터베이스 유형별 기본값", + templateDefaultsEmpty: "아직 어떤 데이터베이스 유형의 기본값으로도 설정되지 않았습니다. 설정하면 해당 데이터베이스 연결에서 AI 패널을 열 때 자동으로 선택됩니다.", + templateDefaultBadgeTitle: "{type}의 기본 템플릿", exportMarkdown: "Markdown으로 내보내기", analysis: "AI 분석", }, diff --git a/apps/desktop/src/i18n/locales/pt-BR.ts b/apps/desktop/src/i18n/locales/pt-BR.ts index f3b537adee..bb639e5d0e 100644 --- a/apps/desktop/src/i18n/locales/pt-BR.ts +++ b/apps/desktop/src/i18n/locales/pt-BR.ts @@ -2766,6 +2766,10 @@ export default withEnglishFallback({ templateSelectorLoading: "Carregando...", templateSelectorEmpty: "Nenhum modelo. Clique em Gerenciar para adicionar.", templateSelectorTooLong: "Os modelos selecionados excedem o limite total de conteúdo ({max} caracteres). Desmarque alguns para continuar.", + templateSetDefault: "Definir como padrão", + templateDefaultsTitle: "Padrão por tipo de banco de dados", + templateDefaultsEmpty: "Ainda não é o padrão de nenhum tipo de banco de dados. Os modelos padrão são aplicados automaticamente ao abrir o painel de IA em uma conexão correspondente.", + templateDefaultBadgeTitle: "Modelo padrão para {type}", exportMarkdown: "Exportar como Markdown", analysis: "AI Analysis", agentSteps: { diff --git a/apps/desktop/src/i18n/locales/zh-CN.ts b/apps/desktop/src/i18n/locales/zh-CN.ts index 91aa14d415..ac276e74f2 100644 --- a/apps/desktop/src/i18n/locales/zh-CN.ts +++ b/apps/desktop/src/i18n/locales/zh-CN.ts @@ -2868,6 +2868,10 @@ export default withEnglishFallback({ templateSelectorLoading: "加载中...", templateSelectorEmpty: "暂无模板，请点击管理添加。", templateSelectorTooLong: "已选模板内容合计超过限制（最多 {max} 字符），请取消部分勾选。", + templateSetDefault: "设为默认", + templateDefaultsTitle: "按数据库类型设为默认", + templateDefaultsEmpty: "尚未设为任何数据库类型的默认模板。设置后，打开对应数据库的 AI 助手时会自动勾选。", + templateDefaultBadgeTitle: "{type} 的默认模板", exportMarkdown: "导出为 Markdown", analysis: "AI 分析", }, diff --git a/apps/desktop/src/i18n/locales/zh-TW.ts b/apps/desktop/src/i18n/locales/zh-TW.ts index 3c6fd9bc0e..a2b291566b 100644 --- a/apps/desktop/src/i18n/locales/zh-TW.ts +++ b/apps/desktop/src/i18n/locales/zh-TW.ts @@ -2657,6 +2657,10 @@ export default withEnglishFallback({ templateSelectorLoading: "載入中...", templateSelectorEmpty: "尚無範本，請點擊管理以新增。", templateSelectorTooLong: "已選範本內容合計超過限制（最多 {max} 字元），請取消部分選取。", + templateSetDefault: "設為預設", + templateDefaultsTitle: "按資料庫類型設為預設", + templateDefaultsEmpty: "尚未設為任何資料庫類型的預設範本。設定後，開啟對應資料庫的 AI 助手時會自動勾選。", + templateDefaultBadgeTitle: "{type} 的預設範本", exportMarkdown: "匯出為 Markdown", analysis: "AI 分析", agentSteps: { diff --git a/apps/desktop/src/lib/__tests__/ai/promptTemplateDefaults.spec.ts b/apps/desktop/src/lib/__tests__/ai/promptTemplateDefaults.spec.ts new file mode 100644 index 0000000000..12da0ee8c2 --- /dev/null +++ b/apps/desktop/src/lib/__tests__/ai/promptTemplateDefaults.spec.ts @@ -0,0 +1,88 @@ +import { describe, expect, it } from "vitest"; + +import { capTemplateIdsToCharLimit, resolveAutoTemplateIds, resolveDefaultTemplateIds } from "@/lib/ai/promptTemplateDefaults"; +import type { PromptTemplate } from "@/types/promptTemplate"; + +function template(id: string, content: string): PromptTemplate { + return { id, name: id, content, createdAt: "2026-01-01", updatedAt: "2026-01-01" }; +} + +describe("resolveDefaultTemplateIds", () => { + it("returns the configured defaults for the db_type", () => { + expect(resolveDefaultTemplateIds("postgresql", { postgresql: ["tpl-pg"] })).toEqual(["tpl-pg"]); + }); + + it("returns empty when no defaults are configured", () => { + // Namespace-switch contract: with no defaults the switch clears the + // selection; last-used is unreachable from this resolution by design. + expect(resolveDefaultTemplateIds("postgresql", {})).toEqual([]); + expect(resolveDefaultTemplateIds("postgresql", { mysql: ["tpl-my"] })).toEqual([]); + }); + + it("resolves nothing without a connection db_type", () => { + expect(resolveDefaultTemplateIds(undefined, { mysql: ["tpl"] })).toEqual([]); + }); +}); + +describe("resolveAutoTemplateIds", () => { + it("returns explicit defaults for the connection db_type", () => { + expect( + resolveAutoTemplateIds({ + dbType: "postgresql", + defaultTemplatesByDbType: { postgresql: ["tpl-pg"] }, + lastUsedTemplatesByDbType: { postgresql: ["tpl-last"] }, + }), + ).toEqual(["tpl-pg"]); + }); + + it("falls back to the db_type's last-used templates only when no defaults are configured", () => { + expect( + resolveAutoTemplateIds({ + dbType: "mysql", + defaultTemplatesByDbType: { postgresql: ["tpl-pg"] }, + lastUsedTemplatesByDbType: { mysql: ["tpl-last"] }, + }), + ).toEqual(["tpl-last"]); + }); + + it("never leaks another db_type's defaults or last-used list", () => { + expect( + resolveAutoTemplateIds({ + dbType: "mysql", + defaultTemplatesByDbType: { postgresql: ["tpl-pg"] }, + lastUsedTemplatesByDbType: { postgresql: ["tpl-last"] }, + }), + ).toEqual([]); + }); + + it("treats an empty defaults list as unconfigured and still applies last-used", () => { + expect( + resolveAutoTemplateIds({ + dbType: "mysql", + defaultTemplatesByDbType: { mysql: [] }, + lastUsedTemplatesByDbType: { mysql: ["tpl-last"] }, + }), + ).toEqual(["tpl-last"]); + }); + + it("resolves nothing without a connection db_type", () => { + expect(resolveAutoTemplateIds({ defaultTemplatesByDbType: { mysql: ["tpl"] }, lastUsedTemplatesByDbType: {} })).toEqual([]); + }); +}); + +describe("capTemplateIdsToCharLimit", () => { + it("drops ids whose templates no longer exist", () => { + const templates = [template("kept", "abc")]; + expect(capTemplateIdsToCharLimit(["missing", "kept"], templates, 100)).toEqual(["kept"]); + }); + + it("skips a template that alone exceeds the remaining budget but keeps later smaller ones", () => { + const templates = [template("small-a", "abc"), template("huge", "x".repeat(10)), template("small-b", "de")]; + expect(capTemplateIdsToCharLimit(["small-a", "huge", "small-b"], templates, 8)).toEqual(["small-a", "small-b"]); + }); + + it("deduplicates repeated ids without double counting their content", () => { + const templates = [template("a", "abc")]; + expect(capTemplateIdsToCharLimit(["a", "a"], templates, 5)).toEqual(["a"]); + }); +}); diff --git a/apps/desktop/src/lib/ai/ai.ts b/apps/desktop/src/lib/ai/ai.ts index 9a6b9839d7..3fb4b741bd 100644 --- a/apps/desktop/src/lib/ai/ai.ts +++ b/apps/desktop/src/lib/ai/ai.ts @@ -716,7 +716,7 @@ async function loadCandidateSchemas(tab: QueryTab, connection: ConnectionConfig) return [database]; } -function aiDatabaseTypeForConnection(connection: ConnectionConfig): DatabaseType { +export function aiDatabaseTypeForConnection(connection: ConnectionConfig): DatabaseType { return effectiveDatabaseTypeForConnection(connection) ?? connection.db_type; } diff --git a/apps/desktop/src/lib/ai/promptTemplateDefaults.ts b/apps/desktop/src/lib/ai/promptTemplateDefaults.ts new file mode 100644 index 0000000000..459fc9356d --- /dev/null +++ b/apps/desktop/src/lib/ai/promptTemplateDefaults.ts @@ -0,0 +1,58 @@ +import type { PromptTemplate } from "@/types/promptTemplate"; +import { promptTemplateCharacterCount } from "@/types/promptTemplate"; + +export interface AutoTemplateResolutionInput { + dbType?: string; + defaultTemplatesByDbType: Record; + lastUsedTemplatesByDbType: Record; +} + +/** + * Template ids explicitly configured as defaults for the db_type, or empty + * when none are configured. Namespace-switch resolution must stop here: + * falling back to last-used would resurrect selections the user just cleared + * on every connection/database/schema change. + */ +export function resolveDefaultTemplateIds(dbType: string | undefined, defaultTemplatesByDbType: Record): string[] { + if (!dbType) return []; + const defaults = defaultTemplatesByDbType[dbType]; + return defaults && defaults.length > 0 ? [...defaults] : []; +} + +/** + * Template ids to auto-select when a panel opens, keyed by the connection's + * db_type: explicit defaults win over the last-used fallback. Both records are + * scoped per db_type so a template tuned for one database never leaks into a + * namespace of another database type. + */ +export function resolveAutoTemplateIds({ dbType, defaultTemplatesByDbType, lastUsedTemplatesByDbType }: AutoTemplateResolutionInput): string[] { + const defaults = resolveDefaultTemplateIds(dbType, defaultTemplatesByDbType); + if (defaults.length > 0) return defaults; + if (!dbType) return []; + const lastUsed = lastUsedTemplatesByDbType[dbType]; + return lastUsed && lastUsed.length > 0 ? [...lastUsed] : []; +} + +/** + * Keep only ids that still exist among templates, deduplicated, dropping any + * single template whose content would push the combined total past maxTotal + * (same budget the manual multi-select enforces). A later smaller template is + * still applied so one oversized default cannot shadow the remaining defaults. + */ +export function capTemplateIdsToCharLimit(ids: string[], templates: PromptTemplate[], maxTotal: number): string[] { + const byId = new Map(templates.map((template) => [template.id, template])); + const kept: string[] = []; + const seen = new Set(); + let total = 0; + for (const id of ids) { + if (seen.has(id)) continue; + seen.add(id); + const template = byId.get(id); + if (!template) continue; + const size = promptTemplateCharacterCount(template.content); + if (total + size > maxTotal) continue; + total += size; + kept.push(id); + } + return kept; +} diff --git a/apps/desktop/src/stores/__tests__/settingsStore.promptTemplateDefaults.spec.ts b/apps/desktop/src/stores/__tests__/settingsStore.promptTemplateDefaults.spec.ts new file mode 100644 index 0000000000..603975cddc --- /dev/null +++ b/apps/desktop/src/stores/__tests__/settingsStore.promptTemplateDefaults.spec.ts @@ -0,0 +1,119 @@ +import { createPinia, setActivePinia } from "pinia"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const saveAiChatSelection = vi.fn(async () => {}); +const loadAiChatSelection = vi.fn(async () => null); + +vi.mock("@/lib/backend/api", () => ({ + loadAiConfigs: vi.fn(async () => []), + loadAiChatSelection: () => loadAiChatSelection(), + saveAiChatSelection: (selection: unknown) => saveAiChatSelection(selection), + loadAiConfig: vi.fn(async () => null), + loadAiProviderConfigs: vi.fn(async () => null), + saveAiConfig: vi.fn(async () => {}), +})); + +import { useSettingsStore } from "@/stores/settingsStore"; + +async function flushed(): Promise { + await new Promise((resolve) => setTimeout(resolve, 0)); +} + +describe("settingsStore per-db_type prompt template selection", () => { + beforeEach(() => { + setActivePinia(createPinia()); + saveAiChatSelection.mockClear(); + loadAiChatSelection.mockClear(); + loadAiChatSelection.mockResolvedValue(null); + }); + + it("setDefaultTemplatesForDbType stores the ids and persists them in the chat selection", async () => { + const settings = useSettingsStore(); + settings.setDefaultTemplatesForDbType("postgresql", ["tpl-1", " tpl-2 ", "tpl-1"]); + expect(settings.aiDefaultTemplatesByDbType).toEqual({ postgresql: ["tpl-1", "tpl-2"] }); + await flushed(); + expect(saveAiChatSelection).toHaveBeenCalled(); + const payload = saveAiChatSelection.mock.calls.at(-1)?.[0]; + expect(payload.defaultTemplatesByDbType).toEqual({ postgresql: ["tpl-1", "tpl-2"] }); + }); + + it("setDefaultTemplatesForDbType with an empty list clears the db_type entry", async () => { + const settings = useSettingsStore(); + settings.setDefaultTemplatesForDbType("mysql", ["tpl-1"]); + await flushed(); + saveAiChatSelection.mockClear(); + settings.setDefaultTemplatesForDbType("mysql", []); + expect(settings.aiDefaultTemplatesByDbType).toEqual({}); + await flushed(); + expect(saveAiChatSelection).toHaveBeenCalled(); + const payload = saveAiChatSelection.mock.calls.at(-1)?.[0]; + // Empty records are omitted entirely, matching the backend's skip_serializing_if. + expect(payload.defaultTemplatesByDbType).toBeUndefined(); + }); + + it("recordLastUsedTemplates skips persisting an empty send for an unrecorded db_type", async () => { + const settings = useSettingsStore(); + settings.recordLastUsedTemplates("mysql", []); + expect(settings.aiLastUsedTemplatesByDbType).toEqual({}); + await flushed(); + expect(saveAiChatSelection).not.toHaveBeenCalled(); + }); + + it("recordLastUsedTemplates clears the remembered selection when the user sends with no templates", async () => { + // Regression: a deselected-everything send used to leave the stale entry, + // resurrecting the old templates the next time a panel opened. + const settings = useSettingsStore(); + settings.recordLastUsedTemplates("mysql", ["tpl-1"]); + await flushed(); + saveAiChatSelection.mockClear(); + + settings.recordLastUsedTemplates("mysql", []); + expect(settings.aiLastUsedTemplatesByDbType).toEqual({}); + await flushed(); + expect(saveAiChatSelection).toHaveBeenCalled(); + const payload = saveAiChatSelection.mock.calls.at(-1)?.[0]; + // Empty records are omitted entirely, matching the backend's skip_serializing_if. + expect(payload.lastUsedTemplatesByDbType).toBeUndefined(); + }); + + it("removeTemplateFromDefaultAndLastUsed prunes the id everywhere and drops emptied db_types", async () => { + const settings = useSettingsStore(); + settings.setDefaultTemplatesForDbType("postgresql", ["tpl-1", "tpl-2"]); + settings.setDefaultTemplatesForDbType("mysql", ["tpl-1"]); + settings.recordLastUsedTemplates("sqlite", ["tpl-1"]); + await flushed(); + saveAiChatSelection.mockClear(); + + settings.removeTemplateFromDefaultAndLastUsed("tpl-1"); + expect(settings.aiDefaultTemplatesByDbType).toEqual({ postgresql: ["tpl-2"] }); + expect(settings.aiLastUsedTemplatesByDbType).toEqual({}); + await flushed(); + expect(saveAiChatSelection).toHaveBeenCalled(); + }); + + it("removeTemplateFromDefaultAndLastUsed persists nothing for unknown ids", async () => { + const settings = useSettingsStore(); + settings.setDefaultTemplatesForDbType("postgresql", ["tpl-1"]); + await flushed(); + saveAiChatSelection.mockClear(); + + settings.removeTemplateFromDefaultAndLastUsed("missing"); + expect(settings.aiDefaultTemplatesByDbType).toEqual({ postgresql: ["tpl-1"] }); + await flushed(); + expect(saveAiChatSelection).not.toHaveBeenCalled(); + }); + + it("initAiConfigs loads and normalizes both per-db_type records from the saved selection", async () => { + loadAiChatSelection.mockResolvedValue({ + version: 1, + effortPreferences: [], + defaultMode: "ask", + defaultTemplatesByDbType: { postgresql: ["tpl-1", " ", "tpl-1"], broken: "not-an-array" }, + lastUsedTemplatesByDbType: { mysql: ["tpl-9"] }, + }); + const settings = useSettingsStore(); + await settings.initAiConfigs(); + expect(settings.aiDefaultTemplatesByDbType).toEqual({ postgresql: ["tpl-1"] }); + expect(settings.aiLastUsedTemplatesByDbType).toEqual({ mysql: ["tpl-9"] }); + }); +}); diff --git a/apps/desktop/src/stores/settingsStore.ts b/apps/desktop/src/stores/settingsStore.ts index f40d746b62..cfe13464fb 100644 --- a/apps/desktop/src/stores/settingsStore.ts +++ b/apps/desktop/src/stores/settingsStore.ts @@ -1537,6 +1537,18 @@ function editorSettingsPatchSnapshot(settings: Partial): Partial return JSON.parse(JSON.stringify(settings)) as Partial; } +/** Keep only well-formed, non-empty template id lists keyed by db_type. */ +function normalizeTemplateIdsByDbType(value?: Record): Record { + if (!value) return {}; + const out: Record = {}; + for (const [dbType, ids] of Object.entries(value)) { + if (!Array.isArray(ids)) continue; + const cleaned = [...new Set(ids.filter((id): id is string => typeof id === "string" && id.trim() !== "").map((id) => id.trim()))]; + if (cleaned.length > 0) out[dbType] = cleaned; + } + return out; +} + export interface SettingsNavigationRequest { id: number; tab: string; @@ -1549,6 +1561,10 @@ export const useSettingsStore = defineStore("settings", () => { const activeModel = ref<{ configId: string; modelId: string } | null>(null); const effortPreferences = ref([]); const defaultAiMode = ref("ask"); + // Per-db_type prompt template defaults (explicit opt-in) and last-used + // fallback; both resolved when an AI panel mounts or its namespace changes. + const aiDefaultTemplatesByDbType = ref>({}); + const aiLastUsedTemplatesByDbType = ref>({}); const isAiConfigLoaded = ref(false); const aiConfigs = ref([]); const desktopSettings = ref({ ...DEFAULT_DESKTOP_SETTINGS }); @@ -1738,6 +1754,8 @@ export const useSettingsStore = defineStore("settings", () => { const savedSelection = await api.loadAiChatSelection().catch(() => null); effortPreferences.value = (savedSelection?.effortPreferences ?? []).filter((preference) => aiConfigs.value.some((config) => config.id === preference.configId)); defaultAiMode.value = savedSelection?.defaultMode ?? "ask"; + aiDefaultTemplatesByDbType.value = normalizeTemplateIdsByDbType(savedSelection?.defaultTemplatesByDbType); + aiLastUsedTemplatesByDbType.value = normalizeTemplateIdsByDbType(savedSelection?.lastUsedTemplatesByDbType); const savedActive = savedSelection?.active; const savedConfig = savedActive ? aiConfigs.value.find((config) => config.id === savedActive.configId) : undefined; @@ -1876,6 +1894,57 @@ export const useSettingsStore = defineStore("settings", () => { persistAiChatSelection(); } + /** Empty id list clears the db_type entry so unsetting a default is expressible. */ + function setDefaultTemplatesForDbType(dbType: string, templateIds: string[]) { + if (!dbType) return; + const next = { ...aiDefaultTemplatesByDbType.value }; + const cleaned = [...new Set(templateIds.map((id) => id.trim()).filter((id) => id !== ""))]; + if (cleaned.length === 0) delete next[dbType]; + else next[dbType] = cleaned; + aiDefaultTemplatesByDbType.value = next; + persistAiChatSelection(); + } + + /** + * Called on send: a non-empty selection is remembered for the db_type, while + * sending with every template deselected is an explicit choice that clears + * the remembered selection — otherwise the stale entry would resurrect the + * old templates the next time a panel opens. + */ + function recordLastUsedTemplates(dbType: string, templateIds: string[]) { + if (!dbType) return; + if (templateIds.length === 0) { + if (!(dbType in aiLastUsedTemplatesByDbType.value)) return; + const next = { ...aiLastUsedTemplatesByDbType.value }; + delete next[dbType]; + aiLastUsedTemplatesByDbType.value = next; + persistAiChatSelection(); + return; + } + aiLastUsedTemplatesByDbType.value = { ...aiLastUsedTemplatesByDbType.value, [dbType]: [...templateIds] }; + persistAiChatSelection(); + } + + function removeTemplateFromDefaultAndLastUsed(templateId: string) { + const strip = (record: Record): Record => { + let changed = false; + const next: Record = {}; + for (const [dbType, ids] of Object.entries(record)) { + const filtered = ids.filter((id) => id !== templateId); + if (filtered.length !== ids.length) changed = true; + if (filtered.length > 0) next[dbType] = filtered; + } + // Identity-preserving no-op lets the caller skip a needless persist. + return changed ? next : record; + }; + const nextDefaults = strip(aiDefaultTemplatesByDbType.value); + const nextLastUsed = strip(aiLastUsedTemplatesByDbType.value); + if (nextDefaults === aiDefaultTemplatesByDbType.value && nextLastUsed === aiLastUsedTemplatesByDbType.value) return; + aiDefaultTemplatesByDbType.value = nextDefaults; + aiLastUsedTemplatesByDbType.value = nextLastUsed; + persistAiChatSelection(); + } + function persistAiChatSelection() { pendingAiChatSelection = { version: 1, @@ -1885,6 +1954,11 @@ export const useSettingsStore = defineStore("settings", () => { selection: { ...preference.selection }, })), defaultMode: defaultAiMode.value, + // Match the backend's skip_serializing_if(empty): omit the per-db_type + // records entirely while nothing is configured so the payload stays + // identical to the pre-defaults format for users without template picks. + ...(Object.keys(aiDefaultTemplatesByDbType.value).length > 0 ? { defaultTemplatesByDbType: aiDefaultTemplatesByDbType.value } : {}), + ...(Object.keys(aiLastUsedTemplatesByDbType.value).length > 0 ? { lastUsedTemplatesByDbType: aiLastUsedTemplatesByDbType.value } : {}), }; if (!aiChatSelectionSaveRunning) void flushAiChatSelection(); } @@ -2257,6 +2331,11 @@ export const useSettingsStore = defineStore("settings", () => { activeEffort, defaultAiMode, setDefaultAiMode, + aiDefaultTemplatesByDbType, + aiLastUsedTemplatesByDbType, + setDefaultTemplatesForDbType, + recordLastUsedTemplates, + removeTemplateFromDefaultAndLastUsed, isAiConfigLoaded, aiConfigs, initAiConfigs, diff --git a/apps/desktop/src/types/ai.ts b/apps/desktop/src/types/ai.ts index 6a2010172f..89068babe3 100644 --- a/apps/desktop/src/types/ai.ts +++ b/apps/desktop/src/types/ai.ts @@ -92,4 +92,8 @@ export interface AiChatSelectionState { active?: AiActiveModelSelection; effortPreferences: AiModelEffortPreference[]; defaultMode?: AiAssistantMode; + /** Prompt template ids auto-applied when the AI panel opens, keyed by connection db_type. */ + defaultTemplatesByDbType?: Record; + /** Prompt template ids from the most recent send, keyed by connection db_type. */ + lastUsedTemplatesByDbType?: Record; } diff --git a/crates/dbx-core/src/ai.rs b/crates/dbx-core/src/ai.rs index c17c833a7d..f840ce6043 100644 --- a/crates/dbx-core/src/ai.rs +++ b/crates/dbx-core/src/ai.rs @@ -3,7 +3,7 @@ use futures::StreamExt; use reqwest::header::{HeaderMap, HeaderName, HeaderValue, AUTHORIZATION, CONTENT_TYPE}; use serde::{Deserialize, Serialize}; use serde_json::json; -use std::collections::{HashMap, HashSet}; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::net::IpAddr; use std::path::Path; use std::sync::{Arc, LazyLock}; @@ -313,6 +313,14 @@ pub struct AiChatSelectionState { pub effort_preferences: Vec, #[serde(default, skip_serializing_if = "Option::is_none")] pub default_mode: Option, + /// Prompt template ids auto-applied when the AI panel opens, keyed by + /// connection db_type. BTreeMap keeps serialized key order stable. + #[serde(default, skip_serializing_if = "BTreeMap::is_empty")] + pub default_templates_by_db_type: BTreeMap>, + /// Prompt template ids from the most recent send, keyed by connection + /// db_type; used as fallback when no defaults are configured. + #[serde(default, skip_serializing_if = "BTreeMap::is_empty")] + pub last_used_templates_by_db_type: BTreeMap>, } impl Default for AiChatSelectionState { @@ -322,6 +330,8 @@ impl Default for AiChatSelectionState { active: None, effort_preferences: Vec::new(), default_mode: None, + default_templates_by_db_type: BTreeMap::new(), + last_used_templates_by_db_type: BTreeMap::new(), } } } diff --git a/crates/dbx-core/src/storage.rs b/crates/dbx-core/src/storage.rs index 586de2da0d..822ad06534 100644 --- a/crates/dbx-core/src/storage.rs +++ b/crates/dbx-core/src/storage.rs @@ -5113,6 +5113,7 @@ mod tests { }; use crate::saved_sql::SavedSqlFile; use rusqlite::{Connection, TransactionBehavior}; + use std::collections::BTreeMap; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; fn temp_db_path(name: &str) -> std::path::PathBuf { @@ -7242,6 +7243,8 @@ mod tests { selection: AiEffortSelection::Enum("high".to_string()), }], default_mode: Some(AiAssistantMode::Agent), + default_templates_by_db_type: BTreeMap::from([("postgresql".to_string(), vec!["tpl-1".to_string()])]), + last_used_templates_by_db_type: BTreeMap::from([("mysql".to_string(), vec!["tpl-2".to_string()])]), }; storage.save_ai_chat_selection(&selection).await.unwrap(); @@ -7250,6 +7253,39 @@ mod tests { assert_eq!(storage.load_app_settings_json().await.unwrap().get("ai_chat_selection_v1"), None); } + // Selection JSON written before per-db-type prompt template defaults existed + // must still deserialize; the new maps fall back to empty. + #[tokio::test] + async fn ai_chat_selection_loads_legacy_payload_without_template_defaults() { + let path = temp_db_path("ai-chat-selection-legacy"); + let storage = Storage::open(&path).await.unwrap(); + let legacy = serde_json::json!({ + "version": 1, + "active": { "configId": "config-1", "modelId": "model-1" }, + "effortPreferences": [], + "defaultMode": "ask" + }); + storage.save_app_state_value(super::APP_STATE_AI_CHAT_SELECTION_KEY, &legacy).await.unwrap(); + + let loaded = storage.load_ai_chat_selection().await.unwrap().unwrap(); + assert_eq!( + loaded.active, + Some(AiActiveModelSelection { config_id: "config-1".to_string(), model_id: "model-1".to_string() }) + ); + assert!(loaded.default_templates_by_db_type.is_empty()); + assert!(loaded.last_used_templates_by_db_type.is_empty()); + } + + // Serialization must omit the per-db-type maps while empty so the payload + // stays identical to the pre-defaults format for users without picks. + #[test] + fn ai_chat_selection_serialization_omits_empty_template_maps() { + let json = serde_json::to_value(AiChatSelectionState::default()).unwrap(); + let object = json.as_object().unwrap(); + assert!(!object.contains_key("defaultTemplatesByDbType")); + assert!(!object.contains_key("lastUsedTemplatesByDbType")); + } + #[tokio::test] async fn tab_runtime_cache_roundtrips_binary_payloads() { let path = temp_db_path("tab-runtime-cache"); diff --git a/docs/content/docs/ai-assistant.mdx b/docs/content/docs/ai-assistant.mdx index 67c8de6564..6e78e7ee42 100644 --- a/docs/content/docs/ai-assistant.mdx +++ b/docs/content/docs/ai-assistant.mdx @@ -112,6 +112,8 @@ Settings can store: Multiple templates can be active at once. DBX snapshots the selected templates and global instructions when the request is sent, so later edits do not alter an already running task. +Templates can also be marked as the **default for one or more database types** (star icon next to a template in Settings → AI). When the AI panel opens on a connection of that database type — or after switching to a different connection, database, or schema — the matching defaults are applied automatically, subject to the same total content limit as manual selection. When no defaults are configured for the current database type, a newly opened panel restores the templates last used with that database type; sending with no templates selected clears that remembered selection. Switching connection, database, or schema applies only explicit defaults — with none configured, the selection starts empty rather than restored. Manual deselection always wins until the next panel open or namespace switch. + Templates enter model context. Do not put passwords, tokens, or unnecessary business data in them. Keep rules concise and non-conflicting because models may not resolve contradictory instructions as expected. ## Writes and Production Safety