fix(ai): preserve Dameng database and schema scopes

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

@@ -55,7 +55,7 @@ import ConnectionGroupBadge from "@/components/connection/ConnectionGroupBadge.v
 import { useQueryStore } from "@/stores/queryStore";
 import { useToast } from "@/composables/useToast";
 import { useNavigationTargets } from "@/composables/useNavigationTargets";
-import { buildAiContext, runAgentStream, isVectorDbType, isValidActionForMode, defaultActionForMode, type AiAction, type AiAssistantMode, type AiSqlFileContext, type CustomPromptContext } from "@/lib/ai/ai";
+import { buildAiContext, resolveAiDatabaseTarget, resolveAiNamespaceSelection, resolveDefaultAiSchema, runAgentStream, isVectorDbType, isValidActionForMode, defaultActionForMode, type AiAction, type AiAssistantMode, type AiSqlFileContext, type CustomPromptContext } from "@/lib/ai/ai";
 import { isAiConfigModelCandidate } from "@/lib/ai/aiConfigCandidates";
 import { orderAiConfigsForDisplay } from "@/lib/ai/aiConfigOrdering";
 import { effortSelectionEquals, runtimeEffortFromPreference } from "@/lib/ai/aiEffortPreference";
@@ -75,7 +75,7 @@ import { aiCancelStream, saveAiConversation, loadAiConversations, deleteAiConver
 import type { AiMessage } from "@/lib/backend/api";
 import type { AiConfigItem, AiEffortCapability, AiEffortOption, AiEffortSelection } from "@/types/ai";
 import type { ConnectionConfig, QueryTab, SavedSqlFile, TableInfo } from "@/types/database";
-import { useDatabaseOptions } from "@/composables/useDatabaseOptions";
+import { fetchNamespaceOptionsForConnection, useDatabaseOptions } from "@/composables/useDatabaseOptions";
 import { decodeSelectableDatabaseValue, encodeSelectableDatabaseValue, formatDatabaseLabel, resolveDefaultDatabase } from "@/lib/database/defaultDatabase";
 import { normalizeSqliteNamespace } from "@/lib/database/sqliteNamespace";
 import { isSchemaAware } from "@/lib/database/databaseCapabilities";
@@ -730,6 +730,7 @@ let confirmedWriteSqlText: string | undefined = undefined;
  *  to prevent a database change between confirmation and execution. */
 let confirmedConnectionId: string | undefined = undefined;
 let confirmedDatabase: string | undefined = undefined;
+let confirmedSchema: string | undefined = undefined;
 
 /** Clear all pending write-confirmation state. Call on every early-return
  *  and failure path so a stale grant cannot leak into a subsequent send(). */
@@ -738,9 +739,13 @@ function clearPendingWriteGrant() {
   confirmedWriteSqlText = undefined;
   confirmedConnectionId = undefined;
   confirmedDatabase = undefined;
+  confirmedSchema = undefined;
 }
 
-const productionContext = computed(() => productionContextForDatabase(props.connection, props.tab?.database));
+const productionContext = computed(() => {
+  const target = props.connection && props.tab ? resolveAiDatabaseTarget(props.tab, props.connection) : undefined;
+  return productionContextForDatabase(props.connection, target?.database);
+});
 
 function sendProposalReply(positive: boolean) {
   // Disable while a stream is in flight or no proposal is currently active.
@@ -762,7 +767,11 @@ function sendProposalReply(positive: boolean) {
     if (confirmedWriteSqlText) {
       allowWriteSqlForNextRun = true;
       confirmedConnectionId = props.connection?.id;
-      confirmedDatabase = props.tab?.database || "";
+      if (props.tab && props.connection) {
+        const target = resolveAiDatabaseTarget(props.tab, props.connection);
+        confirmedDatabase = target.database;
+        confirmedSchema = target.schema;
+      }
     }
     // When no SQL code block is found in the proposal, treat the
     // confirmation as rejected — we cannot bind the agent to a
@@ -804,11 +813,17 @@ function selectModeActionItem(action: AiAction) {
   modeActionOpen.value = false;
 }
 
-const { databaseOptions: allDbOptions, loadDatabaseOptions } = useDatabaseOptions();
+const { databaseOptions, loadDatabaseOptions } = useDatabaseOptions();
+
+// Dameng presents schemas as its top-level namespace, unlike the other
+// connection types that rely on the shared database-options loader.
+const aiDatabaseOptions = ref<Record<string, string[]>>({});
 
 const dbOptions = computed(() => {
-  if (!props.connection) return [];
-  return allDbOptions.value[props.connection.id] || [];
+  const connection = props.connection;
+  if (!connection) return [];
+  if (connection.db_type === "dameng") return aiDatabaseOptions.value[connection.id] || [];
+  return databaseOptions.value[connection.id] || [];
 });
 
 const dbSelectOptions = computed(() => {
@@ -824,20 +839,29 @@ const dbSelectOptions = computed(() => {
   }));
 });
 
-const selectedDatabaseSelectValue = computed(() => (props.connection ? encodeSelectableDatabaseValue(props.connection.db_type, props.tab?.database || "") : ""));
+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) : ""));
 
 const selectedDatabaseLabel = computed(() => {
   if (!props.connection) return t("editor.selectDatabase");
   if (!props.tab) return t("editor.selectDatabase");
-  return formatDatabaseLabel(props.connection, props.tab.database || "", {
+  return formatDatabaseLabel(props.connection, selectedNamespace.value, {
     defaultDatabase: t("editor.defaultDatabase"),
     noDatabase: t("editor.noDatabase"),
   });
 });
 
-async function loadDatabases() {
-  if (!props.connection) return;
-  await loadDatabaseOptions(props.connection.id);
+async function loadDatabases(connection = props.connection): Promise<string[]> {
+  if (!connection) return [];
+  if (connection.db_type !== "dameng") {
+    await loadDatabaseOptions(connection.id);
+    return databaseOptions.value[connection.id] || [];
+  }
+  await connectionStore.ensureConnected(connection.id);
+  const options = await fetchNamespaceOptionsForConnection(connection.id, connection);
+  aiDatabaseOptions.value[connection.id] = options;
+  return options;
 }
 
 async function changeConnection(connectionId: string) {
@@ -845,16 +869,16 @@ async function changeConnection(connectionId: string) {
   if (!conn) return;
   connectionStore.activeConnectionId = connectionId;
   const tab = props.tab;
+  const tabId = tab ? tab.id : queryStore.createTab(connectionId, resolveDefaultDatabase(conn, []));
   if (tab) {
     queryStore.updateConnection(tab.id, connectionId, resolveDefaultDatabase(conn, []));
-  } else {
-    queryStore.createTab(connectionId, resolveDefaultDatabase(conn, []));
   }
   try {
-    await loadDatabaseOptions(connectionId);
-    const database = resolveDefaultDatabase(conn, allDbOptions.value[connectionId] || []);
-    if (tab) {
-      queryStore.updateDatabase(tab.id, database);
+    const options = await loadDatabases(conn);
+    if (conn.db_type === "dameng") {
+      queryStore.updateSchema(tabId, resolveDefaultAiSchema(conn, options));
+    } else {
+      queryStore.updateDatabase(tabId, resolveDefaultDatabase(conn, options));
     }
   } catch (e: unknown) {
     const message = e instanceof Error ? e.message : String(e);
@@ -862,11 +886,16 @@ async function changeConnection(connectionId: string) {
   }
 }
 
-function changeDatabase(value: string) {
+function changeNamespace(value: string) {
   const tab = props.tab;
   const connection = props.connection;
   if (!tab || !connection) return;
-  queryStore.updateDatabase(tab.id, decodeSelectableDatabaseValue(connection.db_type, value));
+  const namespace = decodeSelectableDatabaseValue(connection.db_type, value);
+  if (resolveAiNamespaceSelection(tab, connection).kind === "schema") {
+    queryStore.updateSchema(tab.id, namespace || undefined);
+  } else {
+    queryStore.updateDatabase(tab.id, namespace);
+  }
 }
 
 function flushAssistantDeltas() {
@@ -1717,7 +1746,9 @@ async function send() {
         if (msg.role === "assistant" && msg.content) {
           confirmedWriteSqlText = extractSingleSqlCodeBlock(msg.content);
           confirmedConnectionId = connection.id;
-          confirmedDatabase = tab.database || "";
+          const target = resolveAiDatabaseTarget(tab, connection);
+          confirmedDatabase = target.database;
+          confirmedSchema = target.schema;
           break;
         }
         if (msg.role === "user") break;
@@ -1727,11 +1758,12 @@ async function send() {
       }
     }
   }
-  // Verify the connection/database haven't changed since the user confirmed
-  // the write operation. If the user switched connections or databases between
+  // Verify the connection/database/schema haven't changed since the user confirmed
+  // the write operation. If the user switched connections or namespaces between
   // confirmation and execution, the grant is void.
   if (allowWriteSqlForNextRun && confirmedWriteSqlText) {
-    if (confirmedConnectionId !== connection.id || confirmedDatabase !== (tab.database || "")) {
+    const target = resolveAiDatabaseTarget(tab, connection);
+    if (confirmedConnectionId !== connection.id || confirmedDatabase !== target.database || confirmedSchema !== target.schema) {
       allowWriteSqlForNextRun = false;
       confirmedWriteSqlText = undefined;
     }
@@ -1743,10 +1775,12 @@ async function send() {
   // state, so the values survive to be passed through to the backend.
   const confirmedTargetConnId = allowWriteSql ? confirmedConnectionId : undefined;
   const confirmedTargetDb = allowWriteSql ? confirmedDatabase : undefined;
+  const confirmedTargetSchema = allowWriteSql ? confirmedSchema : undefined;
   allowWriteSqlForNextRun = false;
   confirmedWriteSqlText = undefined;
   confirmedConnectionId = undefined;
   confirmedDatabase = undefined;
+  confirmedSchema = undefined;
   messages.value.push({ role: "assistant", content: "", sourceConnectionName: connection.name });
   const assistantIdx = messages.value.length - 1;
   const sessionId = uuid();
@@ -1770,6 +1804,7 @@ async function send() {
         confirmedWriteSql,
         confirmedConnectionId: confirmedTargetConnId,
         confirmedDatabase: confirmedTargetDb,
+        confirmedSchema: confirmedTargetSchema,
       },
       history,
       (event: AgentEvent) => {
@@ -2398,7 +2433,7 @@ async function openExternalUrl(url: string) {
                   :model-value="selectedDatabaseSelectValue"
                   @update:model-value="
                     (v) => {
-                      if (typeof v === 'string') changeDatabase(v);
+                      if (typeof v === 'string') changeNamespace(v);
                     }
                   "
                   @update:open="