From ccb3667a0c2dece6a90fc0b564c460dd1eb45fff Mon Sep 17 00:00:00 2001
From: t8y2 <1156263951@qq.com>
Date: Wed, 29 Apr 2026 21:32:35 +0800
Subject: [PATCH] feat: Enhance AI Assistant with new actions and context
 handling

---
 src-tauri/src/commands/ai.rs          | 174 ++++++++++++++++++
 src-tauri/src/commands/mod.rs         |   1 +
 src-tauri/src/lib.rs                  |   3 +
 src/App.vue                           | 215 ++++++++++++++++++++--
 src/components/editor/AiAssistant.vue | 254 ++++++++++++++++++++++----
 src/components/editor/QueryEditor.vue |  88 ++++++++-
 src/components/grid/DataGrid.vue      |  25 ++-
 src/composables/useToast.ts           |  16 ++
 src/i18n/locales/en.ts                |  33 ++++
 src/i18n/locales/zh-CN.ts             |  33 ++++
 src/lib/ai.ts                         | 232 ++++++++++++++++++-----
 src/lib/tauri.ts                      |  26 +++
 src/stores/queryStore.ts              |  19 ++
 src/stores/settingsStore.ts           |  32 +++-
 14 files changed, 1045 insertions(+), 106 deletions(-)
 create mode 100644 src-tauri/src/commands/ai.rs
 create mode 100644 src/composables/useToast.ts

diff --git a/src-tauri/src/commands/ai.rs b/src-tauri/src/commands/ai.rs
new file mode 100644
index 000000000..e12e20e79
--- /dev/null
+++ b/src-tauri/src/commands/ai.rs
@@ -0,0 +1,174 @@
+use reqwest::header::{HeaderMap, HeaderValue, AUTHORIZATION, CONTENT_TYPE};
+use serde::{Deserialize, Serialize};
+use serde_json::json;
+use tauri::{AppHandle, Manager};
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "lowercase")]
+pub enum AiProvider {
+    Claude,
+    Openai,
+    Custom,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct AiConfig {
+    pub provider: AiProvider,
+    pub api_key: String,
+    pub endpoint: String,
+    pub model: String,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct AiMessage {
+    pub role: String,
+    pub content: String,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct AiCompletionRequest {
+    pub config: AiConfig,
+    pub system_prompt: String,
+    pub messages: Vec<AiMessage>,
+    pub max_tokens: Option<u32>,
+    pub temperature: Option<f32>,
+}
+
+fn ai_config_file(app: &AppHandle) -> Result<std::path::PathBuf, String> {
+    let dir = app.path().app_data_dir().map_err(|e| e.to_string())?;
+    std::fs::create_dir_all(&dir).map_err(|e| e.to_string())?;
+    Ok(dir.join("ai_config.json"))
+}
+
+#[tauri::command]
+pub async fn save_ai_config(app: AppHandle, config: AiConfig) -> Result<(), String> {
+    let json = serde_json::to_string_pretty(&config).map_err(|e| e.to_string())?;
+    std::fs::write(ai_config_file(&app)?, json).map_err(|e| e.to_string())
+}
+
+#[tauri::command]
+pub async fn load_ai_config(app: AppHandle) -> Result<Option<AiConfig>, String> {
+    let path = ai_config_file(&app)?;
+    if !path.exists() {
+        return Ok(None);
+    }
+    let json = std::fs::read_to_string(path).map_err(|e| e.to_string())?;
+    serde_json::from_str(&json).map(Some).map_err(|e| e.to_string())
+}
+
+#[tauri::command]
+pub async fn ai_complete(request: AiCompletionRequest) -> Result<String, String> {
+    if request.config.api_key.trim().is_empty() {
+        return Err("API key is required".to_string());
+    }
+    if request.config.endpoint.trim().is_empty() {
+        return Err("Endpoint is required".to_string());
+    }
+    if request.config.model.trim().is_empty() {
+        return Err("Model is required".to_string());
+    }
+
+    let client = reqwest::Client::builder()
+        .timeout(std::time::Duration::from_secs(60))
+        .build()
+        .map_err(|e| e.to_string())?;
+
+    match request.config.provider {
+        AiProvider::Claude => call_claude(&client, request).await,
+        AiProvider::Openai | AiProvider::Custom => call_openai_compatible(&client, request).await,
+    }
+}
+
+async fn call_claude(client: &reqwest::Client, request: AiCompletionRequest) -> Result<String, String> {
+    let mut headers = HeaderMap::new();
+    headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
+    headers.insert(
+        "x-api-key",
+        HeaderValue::from_str(&request.config.api_key).map_err(|e| e.to_string())?,
+    );
+    headers.insert("anthropic-version", HeaderValue::from_static("2023-06-01"));
+
+    let body = json!({
+        "model": request.config.model,
+        "max_tokens": request.max_tokens.unwrap_or(2048),
+        "temperature": request.temperature.unwrap_or(0.2),
+        "system": request.system_prompt,
+        "messages": request.messages,
+    });
+
+    let res = client
+        .post(&request.config.endpoint)
+        .headers(headers)
+        .json(&body)
+        .send()
+        .await
+        .map_err(|e| format!("Claude request failed: {e}"))?;
+
+    let status = res.status();
+    let data: serde_json::Value = res.json().await.map_err(|e| e.to_string())?;
+    if !status.is_success() {
+        return Err(extract_error(&data).unwrap_or_else(|| format!("Claude API error: {status}")));
+    }
+
+    Ok(data["content"]
+        .as_array()
+        .and_then(|items| items.iter().find_map(|item| item["text"].as_str()))
+        .unwrap_or_default()
+        .to_string())
+}
+
+async fn call_openai_compatible(
+    client: &reqwest::Client,
+    request: AiCompletionRequest,
+) -> Result<String, String> {
+    let mut headers = HeaderMap::new();
+    headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
+    headers.insert(
+        AUTHORIZATION,
+        HeaderValue::from_str(&format!("Bearer {}", request.config.api_key)).map_err(|e| e.to_string())?,
+    );
+
+    let mut messages = vec![json!({ "role": "system", "content": request.system_prompt })];
+    messages.extend(
+        request
+            .messages
+            .iter()
+            .map(|message| json!({ "role": message.role, "content": message.content })),
+    );
+
+    let body = json!({
+        "model": request.config.model,
+        "messages": messages,
+        "max_tokens": request.max_tokens.unwrap_or(2048),
+        "temperature": request.temperature.unwrap_or(0.2),
+    });
+
+    let res = client
+        .post(&request.config.endpoint)
+        .headers(headers)
+        .json(&body)
+        .send()
+        .await
+        .map_err(|e| format!("AI request failed: {e}"))?;
+
+    let status = res.status();
+    let data: serde_json::Value = res.json().await.map_err(|e| e.to_string())?;
+    if !status.is_success() {
+        return Err(extract_error(&data).unwrap_or_else(|| format!("API error: {status}")));
+    }
+
+    Ok(data["choices"][0]["message"]["content"]
+        .as_str()
+        .unwrap_or_default()
+        .to_string())
+}
+
+fn extract_error(data: &serde_json::Value) -> Option<String> {
+    data["error"]["message"]
+        .as_str()
+        .or_else(|| data["error"].as_str())
+        .map(ToString::to_string)
+}
diff --git a/src-tauri/src/commands/mod.rs b/src-tauri/src/commands/mod.rs
index a38a9b21f..82aa9206d 100644
--- a/src-tauri/src/commands/mod.rs
+++ b/src-tauri/src/commands/mod.rs
@@ -1,3 +1,4 @@
+pub mod ai;
 pub mod connection;
 pub mod history;
 pub mod mongo_cmd;
diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs
index 3435fa6f0..092454c7e 100644
--- a/src-tauri/src/lib.rs
+++ b/src-tauri/src/lib.rs
@@ -31,6 +31,9 @@ pub fn run() {
             }
         })
         .invoke_handler(tauri::generate_handler![
+            commands::ai::ai_complete,
+            commands::ai::save_ai_config,
+            commands::ai::load_ai_config,
             commands::connection::test_connection,
             commands::connection::connect_db,
             commands::connection::disconnect_db,
diff --git a/src/App.vue b/src/App.vue
index e5d96ebc3..49b9c43ac 100644
--- a/src/App.vue
+++ b/src/App.vue
@@ -1,11 +1,17 @@
 <script setup lang="ts">
 import { ref, computed, watch, onMounted, onUnmounted } from "vue";
 import { useI18n } from "vue-i18n";
-import { DatabaseZap, FilePlus2, Play, Loader2, X, Globe, Moon, Sun, Upload, Download, Plus, History } from "lucide-vue-next";
+import { DatabaseZap, FilePlus2, Play, Loader2, X, Globe, Moon, Sun, Upload, Download, Plus, History, Server, Table2 } from "lucide-vue-next";
 import { Splitpanes, Pane } from "splitpanes";
 import "splitpanes/dist/splitpanes.css";
 import { Button } from "@/components/ui/button";
-import { Separator } from "@/components/ui/separator";
+import {
+  Select,
+  SelectContent,
+  SelectItem,
+  SelectTrigger,
+  SelectValue,
+} from "@/components/ui/select";
 import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip";
 import ConnectionTree from "@/components/sidebar/ConnectionTree.vue";
 import ConnectionDialog from "@/components/connection/ConnectionDialog.vue";
@@ -19,6 +25,8 @@ import DangerConfirmDialog from "@/components/editor/DangerConfirmDialog.vue";
 import { useConnectionStore } from "@/stores/connectionStore";
 import { useQueryStore } from "@/stores/queryStore";
 import { useHistoryStore } from "@/stores/historyStore";
+import { useSettingsStore } from "@/stores/settingsStore";
+import { useToast } from "@/composables/useToast";
 import { setLocale, currentLocale, type Locale } from "@/i18n";
 import { getCurrentWindow, type Theme } from "@tauri-apps/api/window";
 import * as api from "@/lib/tauri";
@@ -27,11 +35,15 @@ const { t } = useI18n();
 const connectionStore = useConnectionStore();
 const queryStore = useQueryStore();
 const historyStore = useHistoryStore();
+const settingsStore = useSettingsStore();
+const { message: toastMessage, visible: toastVisible } = useToast();
 
 const showConnectionDialog = ref(false);
 const showHistory = ref(false);
 const dangerSql = ref("");
 const showDangerDialog = ref(false);
+const databaseOptions = ref<Record<string, string[]>>({});
+const loadingDatabaseOptions = ref<Record<string, boolean>>({});
 
 const editConfig = computed(() => {
   const id = connectionStore.editingConnectionId;
@@ -51,6 +63,82 @@ const activeTab = computed(() =>
   queryStore.tabs.find((t) => t.id === queryStore.activeTabId)
 );
 
+const activeConnection = computed(() => {
+  const tab = activeTab.value;
+  return tab ? connectionStore.getConfig(tab.connectionId) : undefined;
+});
+
+const activeTabContext = computed(() => {
+  const tab = activeTab.value;
+  const connection = activeConnection.value;
+  if (!tab || !connection) return [];
+
+  const items = [
+    connection.name,
+    connection.db_type.toUpperCase(),
+  ];
+
+  if (tab.tableMeta?.tableName) {
+    items.push(tab.tableMeta.schema
+      ? `${tab.tableMeta.schema}.${tab.tableMeta.tableName}`
+      : tab.tableMeta.tableName);
+  }
+
+  return items;
+});
+
+const activeDatabaseOptions = computed(() => {
+  const connection = activeConnection.value;
+  return connection ? databaseOptions.value[connection.id] ?? [] : [];
+});
+
+const activeDatabaseValue = computed(() => activeTab.value?.database || "");
+const activeConnectionValue = computed(() => activeConnection.value?.id || "");
+
+function connectionDisplayName(connectionId: string): string {
+  return connectionStore.getConfig(connectionId)?.name || connectionId;
+}
+
+function databaseDisplayName(database: string): string {
+  const connection = activeConnection.value;
+  if (connection?.db_type === "redis" && database !== "") return `db${database}`;
+  return database || t("editor.noDatabase");
+}
+
+async function loadDatabaseOptions(connectionId: string) {
+  const connection = connectionStore.getConfig(connectionId);
+  if (!connection || loadingDatabaseOptions.value[connectionId]) return;
+
+  loadingDatabaseOptions.value[connectionId] = true;
+  try {
+    await connectionStore.ensureConnected(connectionId);
+    if (connection.db_type === "redis") {
+      const dbs = await api.redisListDatabases(connectionId);
+      databaseOptions.value[connectionId] = dbs.map(String);
+    } else if (connection.db_type === "mongodb") {
+      databaseOptions.value[connectionId] = await api.mongoListDatabases(connectionId);
+    } else {
+      const dbs = await api.listDatabases(connectionId);
+      databaseOptions.value[connectionId] = dbs.map((db) => db.name);
+    }
+  } finally {
+    loadingDatabaseOptions.value[connectionId] = false;
+  }
+}
+
+async function getDatabaseOptions(connectionId: string): Promise<string[]> {
+  if (!databaseOptions.value[connectionId]) {
+    await loadDatabaseOptions(connectionId);
+  }
+  return databaseOptions.value[connectionId] ?? [];
+}
+
+watch(activeConnection, (connection) => {
+  if (connection && !databaseOptions.value[connection.id]) {
+    loadDatabaseOptions(connection.id).catch(() => {});
+  }
+}, { immediate: true });
+
 function onEditorUpdate(val: string) {
   if (queryStore.activeTabId) {
     queryStore.updateSql(queryStore.activeTabId, val);
@@ -112,6 +200,36 @@ function onHistoryRestore(sql: string) {
   }
 }
 
+function replaceActiveSql(sql: string) {
+  const tab = activeTab.value;
+  if (!tab) return;
+  queryStore.updateSql(tab.id, sql);
+}
+
+function appendActiveSql(sql: string) {
+  const tab = activeTab.value;
+  if (!tab) return;
+  const current = tab.sql.trimEnd();
+  queryStore.updateSql(tab.id, current ? `${current}\n\n${sql}` : sql);
+}
+
+function changeActiveDatabase(database: any) {
+  const tab = activeTab.value;
+  if (!tab || typeof database !== "string") return;
+  queryStore.updateDatabase(tab.id, database);
+}
+
+async function changeActiveConnection(connectionId: any) {
+  const tab = activeTab.value;
+  if (!tab || typeof connectionId !== "string") return;
+  const connection = connectionStore.getConfig(connectionId);
+  if (!connection) return;
+  const options = await getDatabaseOptions(connectionId);
+  const database = connection.database || options[0] || "";
+  queryStore.updateConnection(tab.id, connectionId, database);
+  connectionStore.activeConnectionId = connectionId;
+}
+
 async function onExecuteSql(sql: string) {
   const tab = activeTab.value;
   if (!tab) return;
@@ -209,6 +327,7 @@ function handleKeydown(e: KeyboardEvent) {
 onMounted(() => {
   applyTheme();
   connectionStore.initFromDisk();
+  settingsStore.initAiConfig();
   window.addEventListener("keydown", handleKeydown);
 });
 
@@ -231,8 +350,6 @@ onUnmounted(() => {
           <TooltipContent>{{ t('toolbar.newConnection') }}</TooltipContent>
         </Tooltip>
 
-        <Separator orientation="vertical" class="h-5" />
-
         <Tooltip>
           <TooltipTrigger as-child>
             <Button variant="ghost" size="icon" class="h-7 w-7" @click="newQuery" :disabled="!connectionStore.activeConnectionId">
@@ -347,6 +464,63 @@ onUnmounted(() => {
 
           <!-- Editor Panel -->
           <div v-if="activeTab" class="flex flex-col flex-1 min-h-0">
+            <div class="h-8 shrink-0 border-b bg-background/80 px-3 flex items-center gap-2 text-xs text-muted-foreground">
+              <Server class="h-3.5 w-3.5 shrink-0" />
+              <Select
+                :model-value="activeConnectionValue"
+                @update:model-value="changeActiveConnection"
+              >
+                <SelectTrigger class="h-6 w-auto max-w-48 border-0 bg-transparent px-1 text-xs font-medium text-foreground shadow-none focus:ring-0">
+                  <SelectValue :placeholder="t('editor.selectConnection')">
+                    {{ connectionDisplayName(activeConnectionValue) }}
+                  </SelectValue>
+                </SelectTrigger>
+                <SelectContent>
+                  <SelectItem
+                    v-for="connection in connectionStore.connections"
+                    :key="connection.id"
+                    :value="connection.id"
+                  >
+                    {{ connection.name }}
+                  </SelectItem>
+                </SelectContent>
+              </Select>
+              <span class="text-muted-foreground/50">/</span>
+              <span class="shrink-0">{{ activeConnection?.db_type.toUpperCase() }}</span>
+              <span class="text-muted-foreground/50">/</span>
+              <Select
+                :model-value="activeDatabaseValue"
+                @update:model-value="changeActiveDatabase"
+                @update:open="(open: boolean) => { if (open && activeConnection) loadDatabaseOptions(activeConnection.id).catch(() => {}) }"
+              >
+                <SelectTrigger class="h-6 w-auto max-w-56 border-0 bg-transparent px-1 text-xs shadow-none focus:ring-0">
+                  <SelectValue :placeholder="loadingDatabaseOptions[activeConnection?.id || ''] ? t('common.loading') : t('editor.selectDatabase')">
+                    {{ databaseDisplayName(activeDatabaseValue) }}
+                  </SelectValue>
+                </SelectTrigger>
+                <SelectContent>
+                  <SelectItem
+                    v-for="database in activeDatabaseOptions"
+                    :key="database"
+                    :value="database"
+                  >
+                    {{ databaseDisplayName(database) }}
+                  </SelectItem>
+                  <SelectItem v-if="!activeDatabaseOptions.length && activeDatabaseValue" :value="activeDatabaseValue">
+                    {{ databaseDisplayName(activeDatabaseValue) }}
+                  </SelectItem>
+                </SelectContent>
+              </Select>
+              <template v-for="item in activeTabContext.slice(2)" :key="item">
+                <span class="text-muted-foreground/50">/</span>
+                <span class="min-w-0 truncate">{{ item }}</span>
+              </template>
+              <span class="flex-1" />
+              <div v-if="activeTab.tableMeta" class="flex min-w-0 items-center gap-1">
+                <Table2 class="h-3.5 w-3.5 shrink-0" />
+                <span class="truncate">{{ activeTab.tableMeta.columns.length }} {{ t('tree.columns') }}</span>
+              </div>
+            </div>
             <!-- Query mode: editor + results -->
             <template v-if="activeTab.mode === 'query'">
               <Splitpanes horizontal class="flex-1">
@@ -358,16 +532,18 @@ onUnmounted(() => {
                       @update:model-value="onEditorUpdate"
                       @execute="tryExecute()"
                     />
-                    <AiAssistant
-                      table-context=""
-                      @insert-sql="(sql: string) => { queryStore.updateSql(activeTab!.id, sql); }"
-                    />
                   </div>
                 </Pane>
                 <Pane :size="60" :min-size="20">
-                  <div class="h-full">
-                    <DataGrid v-if="activeTab.result" :key="activeTab.id" :result="activeTab.result" :sql="activeTab.sql" />
-                    <div v-else class="h-full flex items-center justify-center text-muted-foreground text-sm">
+                  <div class="h-full flex flex-col">
+                    <AiAssistant
+                      :tab="activeTab"
+                      :connection="activeConnection"
+                      @replace-sql="replaceActiveSql"
+                      @append-sql="appendActiveSql"
+                    />
+                    <DataGrid v-if="activeTab.result" :key="activeTab.id" class="flex-1 min-h-0" :result="activeTab.result" :sql="activeTab.sql" />
+                    <div v-else class="flex-1 min-h-0 flex items-center justify-center text-muted-foreground text-sm">
                       {{ t('editor.pressToExecute') }}
                     </div>
                   </div>
@@ -444,6 +620,23 @@ onUnmounted(() => {
 
       <ConnectionDialog v-model:open="showConnectionDialog" :edit-config="editConfig" />
       <DangerConfirmDialog v-model:open="showDangerDialog" :sql="dangerSql" @confirm="onDangerConfirm" />
+
+      <!-- Global Toast -->
+      <Transition name="toast">
+        <div v-if="toastVisible" class="fixed bottom-6 left-1/2 -translate-x-1/2 z-50 px-4 py-2 rounded-lg bg-foreground text-background text-sm shadow-lg">
+          {{ toastMessage }}
+        </div>
+      </Transition>
     </div>
   </TooltipProvider>
 </template>
+
+<style scoped>
+.toast-enter-active, .toast-leave-active {
+  transition: all 0.25s ease;
+}
+.toast-enter-from, .toast-leave-to {
+  opacity: 0;
+  transform: translate(-50%, 8px);
+}
+</style>
diff --git a/src/components/editor/AiAssistant.vue b/src/components/editor/AiAssistant.vue
index f78138156..44114d6f3 100644
--- a/src/components/editor/AiAssistant.vue
+++ b/src/components/editor/AiAssistant.vue
@@ -1,40 +1,100 @@
 <script setup lang="ts">
-import { ref } from "vue";
+import { computed, ref } from "vue";
 import { useI18n } from "vue-i18n";
-import { Sparkles, Loader2, Settings } from "lucide-vue-next";
+import {
+  Bot,
+  Check,
+  Clipboard,
+  Code2,
+  Copy,
+  FilePlus2,
+  Loader2,
+  MessageSquareText,
+  Replace,
+  Settings,
+  Sparkles,
+  Wand2,
+  Wrench,
+} from "lucide-vue-next";
 import { Button } from "@/components/ui/button";
 import { Input } from "@/components/ui/input";
+import { Badge } from "@/components/ui/badge";
 import {
-  Dialog, DialogContent, DialogHeader, DialogTitle, DialogFooter,
+  Dialog,
+  DialogContent,
+  DialogFooter,
+  DialogHeader,
+  DialogTitle,
 } from "@/components/ui/dialog";
 import {
-  Select, SelectContent, SelectItem, SelectTrigger, SelectValue,
+  Select,
+  SelectContent,
+  SelectItem,
+  SelectTrigger,
+  SelectValue,
 } from "@/components/ui/select";
 import { Label } from "@/components/ui/label";
+import { ScrollArea } from "@/components/ui/scroll-area";
 import { useSettingsStore, type AiProvider } from "@/stores/settingsStore";
-import { generateSql } from "@/lib/ai";
+import { buildAiContext, extractSql, runAiAction, type AiAction, type AiContext } from "@/lib/ai";
+import type { ConnectionConfig, QueryTab } from "@/types/database";
 
 const { t } = useI18n();
 const settings = useSettingsStore();
 
 const props = defineProps<{
-  tableContext: string;
+  tab: QueryTab;
+  connection?: ConnectionConfig;
 }>();
 
 const emit = defineEmits<{
-  insertSql: [sql: string];
+  replaceSql: [sql: string];
+  appendSql: [sql: string];
 }>();
 
+const action = ref<AiAction>("generate");
 const prompt = ref("");
 const isGenerating = ref(false);
+const isBuildingContext = ref(false);
 const showSettings = ref(false);
+const showPreview = ref(false);
 const error = ref("");
+const output = ref("");
+const copied = ref(false);
+const lastContext = ref<AiContext | null>(null);
 
 const tempProvider = ref<AiProvider>(settings.aiConfig.provider);
 const tempApiKey = ref(settings.aiConfig.apiKey);
 const tempEndpoint = ref(settings.aiConfig.endpoint);
 const tempModel = ref(settings.aiConfig.model);
 
+const providerDefaults: Record<AiProvider, { endpoint: string; model: string }> = {
+  claude: { endpoint: "https://api.anthropic.com/v1/messages", model: "claude-sonnet-4-20250514" },
+  openai: { endpoint: "https://api.openai.com/v1/chat/completions", model: "gpt-4o" },
+  custom: { endpoint: "", model: "" },
+};
+
+const actionItems: Array<{ value: AiAction; labelKey: string; icon: any }> = [
+  { value: "generate", labelKey: "ai.actions.generate", icon: Sparkles },
+  { value: "explain", labelKey: "ai.actions.explain", icon: MessageSquareText },
+  { value: "optimize", labelKey: "ai.actions.optimize", icon: Wand2 },
+  { value: "fix", labelKey: "ai.actions.fix", icon: Wrench },
+  { value: "convert", labelKey: "ai.actions.convert", icon: Replace },
+  { value: "sampleData", labelKey: "ai.actions.sampleData", icon: FilePlus2 },
+];
+
+const selectedAction = computed(() => actionItems.find((item) => item.value === action.value) ?? actionItems[0]);
+const sqlCandidate = computed(() => extractSql(output.value));
+const canUseSql = computed(() => !!sqlCandidate.value.trim());
+const contextSummary = computed(() => {
+  if (!lastContext.value) return "";
+  const tableCount = lastContext.value.tables.length;
+  return t("ai.contextSummary", {
+    database: lastContext.value.database,
+    tables: tableCount,
+  });
+});
+
 function openSettings() {
   tempProvider.value = settings.aiConfig.provider;
   tempApiKey.value = settings.aiConfig.apiKey;
@@ -53,56 +113,185 @@ function saveSettings() {
   showSettings.value = false;
 }
 
+function selectProvider(provider: AiProvider) {
+  tempProvider.value = provider;
+  tempEndpoint.value = providerDefaults[provider].endpoint;
+  tempModel.value = providerDefaults[provider].model;
+}
+
 async function generate() {
-  if (!prompt.value.trim()) return;
+  if (!props.connection) {
+    error.value = t("ai.noConnection");
+    return;
+  }
   if (!settings.isConfigured()) {
     openSettings();
     return;
   }
+  if (action.value === "generate" && !prompt.value.trim()) return;
+  if (action.value !== "generate" && !props.tab.sql.trim() && !prompt.value.trim()) {
+    error.value = t("ai.noSql");
+    return;
+  }
 
   isGenerating.value = true;
+  isBuildingContext.value = true;
   error.value = "";
+  copied.value = false;
   try {
-    const sql = await generateSql(settings.aiConfig, prompt.value, props.tableContext);
-    emit("insertSql", sql.trim());
-    prompt.value = "";
+    const context = await buildAiContext(props.tab, props.connection);
+    lastContext.value = context;
+    isBuildingContext.value = false;
+    output.value = await runAiAction({
+      config: settings.aiConfig,
+      action: action.value,
+      instruction: prompt.value,
+      context,
+    });
+    showPreview.value = true;
   } catch (e: any) {
     error.value = String(e.message || e);
   } finally {
+    isBuildingContext.value = false;
     isGenerating.value = false;
   }
 }
+
+function replaceSql() {
+  if (!canUseSql.value) return;
+  emit("replaceSql", sqlCandidate.value);
+  showPreview.value = false;
+}
+
+function appendSql() {
+  if (!canUseSql.value) return;
+  emit("appendSql", sqlCandidate.value);
+  showPreview.value = false;
+}
+
+async function copySql() {
+  if (!canUseSql.value) return;
+  await navigator.clipboard.writeText(sqlCandidate.value);
+  copied.value = true;
+  window.setTimeout(() => { copied.value = false; }, 1200);
+}
+
+async function copyAll() {
+  if (!output.value.trim()) return;
+  await navigator.clipboard.writeText(output.value.trim());
+  copied.value = true;
+  window.setTimeout(() => { copied.value = false; }, 1200);
+}
 </script>
 
 <template>
-  <div class="flex items-center gap-1 px-2 py-1 border-t bg-muted/20">
-    <Sparkles class="w-3.5 h-3.5 text-purple-500 shrink-0" />
-    <Input
-      v-model="prompt"
-      class="h-6 text-xs flex-1 border-0 shadow-none focus-visible:ring-0"
-      :placeholder="t('ai.placeholder')"
-      @keydown.enter="generate"
-    />
-    <Button variant="ghost" size="icon" class="h-6 w-6 shrink-0" :disabled="isGenerating" @click="generate">
-      <Loader2 v-if="isGenerating" class="h-3 w-3 animate-spin" />
-      <Sparkles v-else class="h-3 w-3" />
-    </Button>
-    <Button variant="ghost" size="icon" class="h-6 w-6 shrink-0" @click="openSettings">
-      <Settings class="h-3 w-3" />
-    </Button>
-    <span v-if="error" class="text-destructive text-xs truncate max-w-40">{{ error }}</span>
+  <div class="shrink-0 border-b bg-muted/20 px-2 py-1.5">
+    <div class="flex items-center gap-1.5">
+      <Bot class="h-3.5 w-3.5 shrink-0 text-primary" />
+      <Select :model-value="action" @update:model-value="(v: any) => action = v">
+        <SelectTrigger class="h-7 w-32 text-xs">
+          <SelectValue />
+        </SelectTrigger>
+        <SelectContent>
+          <SelectItem v-for="item in actionItems" :key="item.value" :value="item.value">
+            {{ t(item.labelKey) }}
+          </SelectItem>
+        </SelectContent>
+      </Select>
+
+      <Input
+        v-model="prompt"
+        class="h-7 flex-1 border-0 bg-background/70 text-xs shadow-none focus-visible:ring-1"
+        :placeholder="t(`ai.placeholders.${action}`)"
+        @keydown.enter="generate"
+      />
+
+      <Button
+        variant="default"
+        size="xs"
+        class="shrink-0"
+        :disabled="isGenerating"
+        @click="generate"
+      >
+        <Loader2 v-if="isGenerating" class="h-3 w-3 animate-spin" />
+        <component :is="selectedAction.icon" v-else class="h-3 w-3" />
+        <span>{{ isBuildingContext ? t('ai.readingSchema') : t('ai.run') }}</span>
+      </Button>
+
+      <Button variant="ghost" size="icon-xs" class="shrink-0" @click="openSettings">
+        <Settings class="h-3 w-3" />
+      </Button>
+    </div>
+
+    <div v-if="error || contextSummary" class="mt-1 flex items-center gap-2 text-xs">
+      <span v-if="error" class="truncate text-destructive">{{ error }}</span>
+      <span v-else-if="contextSummary" class="truncate text-muted-foreground">{{ contextSummary }}</span>
+      <Badge v-if="lastContext?.truncated" variant="outline" class="h-4 px-1.5 text-[10px]">
+        {{ t('ai.truncated') }}
+      </Badge>
+    </div>
   </div>
 
-  <!-- Settings Dialog -->
+  <Dialog v-model:open="showPreview">
+    <DialogContent class="sm:max-w-3xl">
+      <DialogHeader>
+        <DialogTitle class="flex items-center gap-2">
+          <Code2 class="h-4 w-4" />
+          {{ t(selectedAction.labelKey) }}
+        </DialogTitle>
+      </DialogHeader>
+
+      <div class="grid gap-3">
+        <div v-if="contextSummary" class="flex items-center gap-2 text-xs text-muted-foreground">
+          <span>{{ contextSummary }}</span>
+          <Badge v-if="lastContext?.truncated" variant="outline">{{ t('ai.truncated') }}</Badge>
+        </div>
+
+        <ScrollArea class="h-72 rounded-lg border bg-background">
+          <pre class="whitespace-pre-wrap p-3 text-xs leading-relaxed"><code>{{ output }}</code></pre>
+        </ScrollArea>
+
+        <div v-if="canUseSql" class="rounded-lg border bg-muted/20">
+          <div class="flex items-center justify-between border-b px-3 py-2">
+            <span class="text-xs font-medium">{{ t('ai.sqlPreview') }}</span>
+            <Button variant="ghost" size="xs" @click="copySql">
+              <Check v-if="copied" class="h-3 w-3" />
+              <Copy v-else class="h-3 w-3" />
+              {{ copied ? t('ai.copied') : t('ai.copySql') }}
+            </Button>
+          </div>
+          <ScrollArea class="max-h-48">
+            <pre class="whitespace-pre-wrap p-3 text-xs leading-relaxed"><code>{{ sqlCandidate }}</code></pre>
+          </ScrollArea>
+        </div>
+      </div>
+
+      <DialogFooter class="gap-2 sm:gap-2">
+        <Button variant="outline" size="sm" @click="copyAll">
+          <Clipboard class="h-3.5 w-3.5" />
+          {{ t('ai.copyAll') }}
+        </Button>
+        <Button variant="outline" size="sm" :disabled="!canUseSql" @click="appendSql">
+          <FilePlus2 class="h-3.5 w-3.5" />
+          {{ t('ai.append') }}
+        </Button>
+        <Button size="sm" :disabled="!canUseSql" @click="replaceSql">
+          <Replace class="h-3.5 w-3.5" />
+          {{ t('ai.replace') }}
+        </Button>
+      </DialogFooter>
+    </DialogContent>
+  </Dialog>
+
   <Dialog v-model:open="showSettings">
-    <DialogContent class="sm:max-w-96">
+    <DialogContent class="sm:max-w-[420px]">
       <DialogHeader>
         <DialogTitle>{{ t('ai.settings') }}</DialogTitle>
       </DialogHeader>
-      <div class="grid gap-3 py-3">
+      <div class="grid gap-3 py-2">
         <div class="grid grid-cols-3 items-center gap-3">
           <Label class="text-right text-xs">{{ t('ai.provider') }}</Label>
-          <Select :model-value="tempProvider" @update:model-value="(v: any) => tempProvider = v">
+          <Select :model-value="tempProvider" @update:model-value="(v: any) => selectProvider(v)">
             <SelectTrigger class="col-span-2 h-8 text-xs"><SelectValue /></SelectTrigger>
             <SelectContent>
               <SelectItem value="claude">Claude</SelectItem>
@@ -123,6 +312,9 @@ async function generate() {
           <Label class="text-right text-xs">Model</Label>
           <Input v-model="tempModel" class="col-span-2 h-8 text-xs" />
         </div>
+        <p class="text-xs leading-relaxed text-muted-foreground">
+          {{ t('ai.settingsHint') }}
+        </p>
       </div>
       <DialogFooter>
         <Button size="sm" @click="saveSettings">{{ t('grid.save') }}</Button>
diff --git a/src/components/editor/QueryEditor.vue b/src/components/editor/QueryEditor.vue
index 16dad1613..790efef82 100644
--- a/src/components/editor/QueryEditor.vue
+++ b/src/components/editor/QueryEditor.vue
@@ -14,13 +14,56 @@ const emit = defineEmits<{
 
 const editorRef = ref<HTMLDivElement>();
 const view = shallowRef<EditorViewType | null>(null);
+const DEFAULT_FONT_SIZE = 13;
+const MIN_FONT_SIZE = 10;
+const MAX_FONT_SIZE = 24;
+let editorViewModule: typeof import("@codemirror/view") | null = null;
+let fontSizeTheme: import("@codemirror/state").Compartment | null = null;
+
+const savedFontSize = Number(localStorage.getItem("dbx-query-editor-font-size"));
+const fontSize = ref(
+  Number.isFinite(savedFontSize)
+    ? Math.min(MAX_FONT_SIZE, Math.max(MIN_FONT_SIZE, savedFontSize))
+    : DEFAULT_FONT_SIZE,
+);
+
+function fontTheme(EditorView: typeof import("@codemirror/view").EditorView, size: number) {
+  return EditorView.theme({
+    "&": { height: "100%", fontSize: `${size}px` },
+    ".cm-scroller": { overflow: "auto" },
+    ".cm-content": { fontFamily: "'JetBrains Mono', 'Fira Code', monospace" },
+  });
+}
+
+function setFontSize(size: number) {
+  const next = Math.min(MAX_FONT_SIZE, Math.max(MIN_FONT_SIZE, size));
+  fontSize.value = next;
+  localStorage.setItem("dbx-query-editor-font-size", String(next));
+  if (view.value && fontSizeTheme && editorViewModule) {
+    view.value.dispatch({
+      effects: fontSizeTheme.reconfigure(fontTheme(editorViewModule.EditorView, next)),
+    });
+  }
+}
+
+function zoomIn() {
+  setFontSize(fontSize.value + 1);
+}
+
+function zoomOut() {
+  setFontSize(fontSize.value - 1);
+}
+
+function resetZoom() {
+  setFontSize(DEFAULT_FONT_SIZE);
+}
 
 onMounted(async () => {
   if (!editorRef.value) return;
 
   const [
     { EditorView, keymap },
-    { EditorState },
+    { EditorState, Compartment },
     { sql, MySQL, PostgreSQL },
     { basicSetup },
     { oneDark },
@@ -31,10 +74,40 @@ onMounted(async () => {
     import("codemirror"),
     import("@codemirror/theme-one-dark"),
   ]);
+  editorViewModule = { EditorView, keymap } as typeof import("@codemirror/view");
+  fontSizeTheme = new Compartment();
 
   const dialect = props.dialect === "postgres" ? PostgreSQL : MySQL;
 
   const runKeymap = keymap.of([
+    {
+      key: "Mod-=",
+      run: () => {
+        zoomIn();
+        return true;
+      },
+    },
+    {
+      key: "Mod-+",
+      run: () => {
+        zoomIn();
+        return true;
+      },
+    },
+    {
+      key: "Mod--",
+      run: () => {
+        zoomOut();
+        return true;
+      },
+    },
+    {
+      key: "Mod-0",
+      run: () => {
+        resetZoom();
+        return true;
+      },
+    },
     {
       key: "Mod-Enter",
       run: () => {
@@ -56,10 +129,15 @@ onMounted(async () => {
           emit("update:modelValue", update.state.doc.toString());
         }
       }),
-      EditorView.theme({
-        "&": { height: "100%", fontSize: "13px" },
-        ".cm-scroller": { overflow: "auto" },
-        ".cm-content": { fontFamily: "'JetBrains Mono', 'Fira Code', monospace" },
+      fontSizeTheme.of(fontTheme(EditorView, fontSize.value)),
+      EditorView.domEventHandlers({
+        wheel(event) {
+          if (!event.metaKey && !event.ctrlKey) return false;
+          event.preventDefault();
+          if (event.deltaY < 0) zoomIn();
+          else if (event.deltaY > 0) zoomOut();
+          return true;
+        },
       }),
     ],
   });
diff --git a/src/components/grid/DataGrid.vue b/src/components/grid/DataGrid.vue
index 81cc04987..317675e27 100644
--- a/src/components/grid/DataGrid.vue
+++ b/src/components/grid/DataGrid.vue
@@ -16,7 +16,10 @@ import type { QueryResult, ColumnInfo } from "@/types/database";
 import { save as savePath } from "@tauri-apps/plugin-dialog";
 import { writeTextFile } from "@tauri-apps/plugin-fs";
 
+import { useToast } from "@/composables/useToast";
+
 const { t } = useI18n();
+const { toast } = useToast();
 
 const props = defineProps<{
   result: QueryResult;
@@ -481,7 +484,9 @@ async function exportMarkdown() {
 const sqlOneLiner = computed(() => props.sql?.replace(/\s+/g, " ").trim() || "");
 
 function copySql() {
-  if (props.sql) navigator.clipboard.writeText(props.sql);
+  if (!props.sql) return;
+  navigator.clipboard.writeText(props.sql);
+  toast(t('grid.copied'));
 }
 </script>
 
@@ -509,15 +514,21 @@ function copySql() {
               <div
                 v-for="(col, colIdx) in result.columns"
                 :key="col"
-                class="shrink-0 px-3 py-1.5 border-r border-border whitespace-nowrap cursor-pointer hover:bg-accent/50 select-none relative"
+                class="shrink-0 px-3 py-1.5 border-r border-border whitespace-nowrap cursor-pointer hover:bg-accent/50 select-none relative overflow-hidden"
                 :style="{ width: `var(--col-w-${colIdx})` }"
                 @click="toggleSort(col)"
               >
-                <span class="inline-flex items-center gap-1">
-                  {{ col }}
-                  <ArrowUp v-if="sortCol === col && sortDir === 'asc'" class="w-3 h-3" />
-                  <ArrowDown v-else-if="sortCol === col && sortDir === 'desc'" class="w-3 h-3" />
-                  <span v-if="columnTypeMap.get(col)" class="text-[10px] font-normal ml-auto" :class="typeColorClass(columnTypeMap.get(col)!)">#{{ columnTypeMap.get(col) }}</span>
+                <span class="flex min-w-0 items-center gap-1 overflow-hidden">
+                  <span class="min-w-0 truncate">{{ col }}</span>
+                  <ArrowUp v-if="sortCol === col && sortDir === 'asc'" class="h-3 w-3 shrink-0" />
+                  <ArrowDown v-else-if="sortCol === col && sortDir === 'desc'" class="h-3 w-3 shrink-0" />
+                  <span
+                    v-if="columnTypeMap.get(col)"
+                    class="shrink overflow-hidden truncate text-[10px] font-normal"
+                    :class="typeColorClass(columnTypeMap.get(col)!)"
+                  >
+                    #{{ columnTypeMap.get(col) }}
+                  </span>
                 </span>
                 <div
                   class="absolute right-0 top-0 bottom-0 w-1.5 cursor-col-resize hover:bg-primary/30"
diff --git a/src/composables/useToast.ts b/src/composables/useToast.ts
new file mode 100644
index 000000000..26c18d3e6
--- /dev/null
+++ b/src/composables/useToast.ts
@@ -0,0 +1,16 @@
+import { ref } from "vue";
+
+const message = ref("");
+const visible = ref(false);
+let timer = 0;
+
+export function useToast() {
+  function toast(msg: string, duration = 2000) {
+    message.value = msg;
+    visible.value = true;
+    clearTimeout(timer);
+    timer = window.setTimeout(() => { visible.value = false; }, duration);
+  }
+
+  return { message, visible, toast };
+}
diff --git a/src/i18n/locales/en.ts b/src/i18n/locales/en.ts
index 91b1072fa..46c545dc4 100644
--- a/src/i18n/locales/en.ts
+++ b/src/i18n/locales/en.ts
@@ -42,6 +42,9 @@ export default {
   },
   editor: {
     pressToExecute: "Press Cmd+Enter to execute",
+    noDatabase: "No database selected",
+    selectConnection: "Select connection",
+    selectDatabase: "Select database",
   },
   grid: {
     rows: "{count} rows",
@@ -57,6 +60,7 @@ export default {
     exportCsv: "Export CSV",
     exportJson: "Export JSON",
     exportMarkdown: "Export Markdown",
+    copied: "Copied!",
     search: "Search...",
     page: "Page {page}",
     rowsPerPage: "Rows per page",
@@ -78,6 +82,35 @@ export default {
     placeholder: "Describe your query in natural language...",
     settings: "AI Settings",
     provider: "Provider",
+    run: "Run",
+    readingSchema: "Reading schema",
+    noConnection: "No connection is available for this tab",
+    noSql: "No SQL to process",
+    sqlPreview: "SQL Preview",
+    copySql: "Copy SQL",
+    copyAll: "Copy All",
+    copied: "Copied",
+    replace: "Replace Editor",
+    append: "Append to Editor",
+    truncated: "Context truncated",
+    contextSummary: "{database} · {tables} tables",
+    settingsHint: "The config is stored in the local app data directory. Requests are sent by the Tauri backend instead of directly from the frontend.",
+    actions: {
+      generate: "Generate SQL",
+      explain: "Explain SQL",
+      optimize: "Optimize SQL",
+      fix: "Fix Error",
+      convert: "Convert Dialect",
+      sampleData: "Sample Data",
+    },
+    placeholders: {
+      generate: "Describe what you want to query, e.g. orders per user in the last 7 days",
+      explain: "Optional: add what you want to understand",
+      optimize: "Optional: add a goal, e.g. reduce full table scans",
+      fix: "Paste the error or describe the expected result",
+      convert: "e.g. convert to PostgreSQL / MySQL / SQL Server",
+      sampleData: "Describe the sample data or test statements to generate",
+    },
   },
   contextMenu: {
     openConnection: "Open Connection",
diff --git a/src/i18n/locales/zh-CN.ts b/src/i18n/locales/zh-CN.ts
index a6a6cc0e3..78347970f 100644
--- a/src/i18n/locales/zh-CN.ts
+++ b/src/i18n/locales/zh-CN.ts
@@ -42,6 +42,9 @@ export default {
   },
   editor: {
     pressToExecute: "按 Cmd+Enter 执行查询",
+    noDatabase: "未选择数据库",
+    selectConnection: "选择连接",
+    selectDatabase: "选择数据库",
   },
   grid: {
     rows: "{count} 行",
@@ -57,6 +60,7 @@ export default {
     exportCsv: "导出 CSV",
     exportJson: "导出 JSON",
     exportMarkdown: "导出 Markdown",
+    copied: "已复制!",
     search: "搜索...",
     page: "第 {page} 页",
     rowsPerPage: "每页行数",
@@ -78,6 +82,35 @@ export default {
     placeholder: "用自然语言描述你的查询...",
     settings: "AI 设置",
     provider: "提供商",
+    run: "执行",
+    readingSchema: "读取结构",
+    noConnection: "当前标签页没有可用连接",
+    noSql: "当前没有 SQL 可处理",
+    sqlPreview: "SQL 预览",
+    copySql: "复制 SQL",
+    copyAll: "复制全部",
+    copied: "已复制",
+    replace: "替换编辑器",
+    append: "追加到编辑器",
+    truncated: "上下文已截断",
+    contextSummary: "{database} · {tables} 张表",
+    settingsHint: "配置会保存在本机应用数据目录中。请求由 Tauri 后端发出，避免在前端直接暴露给模型服务。",
+    actions: {
+      generate: "生成 SQL",
+      explain: "解释 SQL",
+      optimize: "优化 SQL",
+      fix: "修复错误",
+      convert: "转换方言",
+      sampleData: "生成样例",
+    },
+    placeholders: {
+      generate: "描述你想查询什么，例如：统计最近 7 天每个用户的订单数",
+      explain: "可留空，或补充你关心的点",
+      optimize: "可留空，或说明优化目标，例如：减少全表扫描",
+      fix: "粘贴报错或说明期望结果",
+      convert: "例如：转换成 PostgreSQL / MySQL / SQL Server",
+      sampleData: "描述要生成的样例数据或测试语句",
+    },
   },
   contextMenu: {
     openConnection: "打开连接",
diff --git a/src/lib/ai.ts b/src/lib/ai.ts
index b9eed63c2..b7ee3c023 100644
--- a/src/lib/ai.ts
+++ b/src/lib/ai.ts
@@ -1,50 +1,196 @@
 import type { AiConfig } from "@/stores/settingsStore";
+import type { ColumnInfo, ConnectionConfig, DatabaseType, QueryResult, QueryTab } from "@/types/database";
+import * as api from "@/lib/tauri";
 
-export async function generateSql(
-  config: AiConfig,
-  prompt: string,
-  tableContext: string,
-): Promise<string> {
-  const systemPrompt = `You are a SQL expert assistant. Given the database schema below, generate a SQL query based on the user's request. Return ONLY the SQL query, no explanations.\n\nSchema:\n${tableContext}`;
+export type AiAction = "generate" | "explain" | "optimize" | "fix" | "convert" | "sampleData";
 
-  if (config.provider === "claude") {
-    const res = await fetch(config.endpoint, {
-      method: "POST",
-      headers: {
-        "Content-Type": "application/json",
-        "x-api-key": config.apiKey,
-        "anthropic-version": "2023-06-01",
-        "anthropic-dangerous-direct-browser-access": "true",
-      },
-      body: JSON.stringify({
-        model: config.model,
-        max_tokens: 2048,
-        system: systemPrompt,
-        messages: [{ role: "user", content: prompt }],
-      }),
+export interface AiSchemaTable {
+  schema?: string;
+  name: string;
+  tableType: string;
+  columns: ColumnInfo[];
+}
+
+export interface AiContext {
+  connectionName: string;
+  databaseType: DatabaseType;
+  database: string;
+  currentSql: string;
+  lastError?: string;
+  lastResultPreview?: string;
+  tables: AiSchemaTable[];
+  truncated: boolean;
+}
+
+export interface AiRequestInput {
+  config: AiConfig;
+  action: AiAction;
+  instruction: string;
+  context: AiContext;
+}
+
+const ACTION_INSTRUCTIONS: Record<AiAction, string> = {
+  generate: "Generate a SQL query that satisfies the user's request.",
+  explain: "Explain the current SQL clearly and point out risky operations or assumptions.",
+  optimize: "Rewrite or suggest improvements for the current SQL. Prefer a complete improved SQL query first, followed by short notes.",
+  fix: "Fix the current SQL using the provided error/result context. Return the corrected SQL first, followed by short notes if needed.",
+  convert: "Convert the current SQL to the target dialect requested by the user. Return the converted SQL first.",
+  sampleData: "Generate safe sample SQL statements or mock data for the current schema. Do not use real production data.",
+};
+
+export async function runAiAction(input: AiRequestInput): Promise<string> {
+  const systemPrompt = buildSystemPrompt(input.action, input.context);
+  const userPrompt = [
+    `Action: ${input.action}`,
+    ACTION_INSTRUCTIONS[input.action],
+    "",
+    "User request:",
+    input.instruction.trim() || "(No extra instruction provided.)",
+  ].join("\n");
+
+  return api.aiComplete({
+    config: input.config,
+    systemPrompt,
+    messages: [{ role: "user", content: userPrompt }],
+    maxTokens: 2400,
+    temperature: 0.15,
+  });
+}
+
+export function extractSql(text: string): string {
+  const fenced = text.match(/```(?:sql|mysql|postgresql|sqlite|tsql|clickhouse)?\s*([\s\S]*?)```/i);
+  if (fenced?.[1]) return fenced[1].trim();
+  return text.trim();
+}
+
+export function buildSystemPrompt(action: AiAction, context: AiContext): string {
+  const schema = formatSchema(context);
+  const resultPreview = context.lastResultPreview
+    ? `\nLast result preview:\n${context.lastResultPreview}\n`
+    : "";
+  const lastError = context.lastError ? `\nLast error:\n${context.lastError}\n` : "";
+
+  return [
+    "You are DBX's built-in database assistant.",
+    "Be precise, conservative, and adapt SQL to the active database dialect.",
+    "Never invent tables or columns that are not present in the schema context unless the user explicitly asks for hypothetical examples.",
+    "For destructive statements such as DROP, DELETE, TRUNCATE, ALTER, or UPDATE without a clear WHERE clause, warn briefly and prefer a safer SELECT preview when appropriate.",
+    "When returning SQL, put the SQL in a fenced ```sql code block. Keep extra explanation short and practical.",
+    action === "generate" ? "For generate actions, return the SQL first and avoid long explanations." : "",
+    "",
+    `Database type: ${context.databaseType}`,
+    `Connection: ${context.connectionName}`,
+    `Database: ${context.database}`,
+    context.truncated ? "Schema context is truncated." : "Schema context is complete within the current budget.",
+    "",
+    `Current SQL:\n${context.currentSql.trim() || "(empty)"}`,
+    lastError,
+    resultPreview,
+    `Schema:\n${schema}`,
+  ].filter(Boolean).join("\n");
+}
+
+function formatSchema(context: AiContext): string {
+  if (!context.tables.length) return "(No table schema loaded.)";
+
+  return context.tables.map((table) => {
+    const name = table.schema ? `${table.schema}.${table.name}` : table.name;
+    const columns = table.columns.map((column) => {
+      const flags = [
+        column.is_primary_key ? "primary key" : "",
+        column.is_nullable ? "nullable" : "not null",
+        column.column_default ? `default ${column.column_default}` : "",
+        column.extra || "",
+      ].filter(Boolean).join(", ");
+      return `  - ${column.name}: ${column.data_type}${flags ? ` (${flags})` : ""}`;
     });
-    if (!res.ok) throw new Error(`Claude API error: ${res.status}`);
-    const data = await res.json();
-    return data.content?.[0]?.text || "";
+    return [`${name} (${table.tableType})`, ...columns].join("\n");
+  }).join("\n\n");
+}
+
+export async function buildAiContext(
+  tab: QueryTab,
+  connection: ConnectionConfig,
+  options: { maxTables?: number; maxColumnsPerTable?: number } = {},
+): Promise<AiContext> {
+  const maxTables = options.maxTables ?? 12;
+  const maxColumnsPerTable = options.maxColumnsPerTable ?? 40;
+  const tables: AiSchemaTable[] = [];
+  let truncated = false;
+
+  if (tab.tableMeta) {
+    tables.push({
+      schema: tab.tableMeta.schema,
+      name: tab.tableMeta.tableName,
+      tableType: "TABLE",
+      columns: tab.tableMeta.columns.slice(0, maxColumnsPerTable),
+    });
+    truncated = tab.tableMeta.columns.length > maxColumnsPerTable;
+  } else if (!["redis", "mongodb"].includes(connection.db_type)) {
+    try {
+      const schemas = await loadCandidateSchemas(tab, connection);
+      for (const schema of schemas) {
+        const tableList = await api.listTables(tab.connectionId, tab.database, schema);
+        for (const table of tableList) {
+          if (tables.length >= maxTables) {
+            truncated = true;
+            break;
+          }
+          const columns = await api.getColumns(tab.connectionId, tab.database, schema, table.name);
+          tables.push({
+            schema: schema === tab.database && connection.db_type !== "postgres" ? undefined : schema,
+            name: table.name,
+            tableType: table.table_type,
+            columns: columns.slice(0, maxColumnsPerTable),
+          });
+          if (columns.length > maxColumnsPerTable) truncated = true;
+        }
+        if (tables.length >= maxTables) break;
+      }
+    } catch {
+      truncated = true;
+    }
   }
 
-  // OpenAI-compatible (works for OpenAI, custom endpoints)
-  const res = await fetch(config.endpoint, {
-    method: "POST",
-    headers: {
-      "Content-Type": "application/json",
-      "Authorization": `Bearer ${config.apiKey}`,
-    },
-    body: JSON.stringify({
-      model: config.model,
-      messages: [
-        { role: "system", content: systemPrompt },
-        { role: "user", content: prompt },
-      ],
-      max_tokens: 2048,
-    }),
-  });
-  if (!res.ok) throw new Error(`API error: ${res.status}`);
-  const data = await res.json();
-  return data.choices?.[0]?.message?.content || "";
+  return {
+    connectionName: connection.name,
+    databaseType: connection.db_type,
+    database: tab.database,
+    currentSql: tab.sql,
+    lastError: extractLastError(tab.result),
+    lastResultPreview: formatResultPreview(tab.result),
+    tables,
+    truncated,
+  };
+}
+
+async function loadCandidateSchemas(tab: QueryTab, connection: ConnectionConfig): Promise<string[]> {
+  if (connection.db_type === "postgres" || connection.db_type === "sqlserver") {
+    const schemas = await api.listSchemas(tab.connectionId, tab.database);
+    return prioritizeSchemas(schemas);
+  }
+  return [tab.database || connection.database || "main"];
+}
+
+function prioritizeSchemas(schemas: string[]): string[] {
+  const preferred = ["public", "dbo", "main"];
+  return [...schemas].sort((a, b) => {
+    const ai = preferred.indexOf(a);
+    const bi = preferred.indexOf(b);
+    if (ai >= 0 || bi >= 0) return (ai >= 0 ? ai : 99) - (bi >= 0 ? bi : 99);
+    return a.localeCompare(b);
+  });
+}
+
+function extractLastError(result?: QueryResult): string | undefined {
+  if (!result?.columns.includes("Error")) return undefined;
+  return result.rows[0]?.[0] == null ? undefined : String(result.rows[0][0]);
+}
+
+function formatResultPreview(result?: QueryResult): string | undefined {
+  if (!result || result.columns.includes("Error") || !result.rows.length) return undefined;
+  const rows = result.rows.slice(0, 5).map((row) => {
+    return result.columns.map((column, index) => `${column}=${JSON.stringify(row[index] ?? null)}`).join(", ");
+  });
+  return rows.join("\n");
 }
diff --git a/src/lib/tauri.ts b/src/lib/tauri.ts
index 6e6b6dc74..e354fc14c 100644
--- a/src/lib/tauri.ts
+++ b/src/lib/tauri.ts
@@ -9,6 +9,32 @@ import type {
   TriggerInfo,
   QueryResult,
 } from "@/types/database";
+import type { AiConfig } from "@/stores/settingsStore";
+
+export interface AiMessage {
+  role: "user" | "assistant" | "system";
+  content: string;
+}
+
+export interface AiCompletionRequest {
+  config: AiConfig;
+  systemPrompt: string;
+  messages: AiMessage[];
+  maxTokens?: number;
+  temperature?: number;
+}
+
+export async function aiComplete(request: AiCompletionRequest): Promise<string> {
+  return invoke("ai_complete", { request });
+}
+
+export async function saveAiConfig(config: AiConfig): Promise<void> {
+  return invoke("save_ai_config", { config });
+}
+
+export async function loadAiConfig(): Promise<AiConfig | null> {
+  return invoke("load_ai_config");
+}
 
 export async function testConnection(config: ConnectionConfig): Promise<string> {
   return invoke("test_connection", { config });
diff --git a/src/stores/queryStore.ts b/src/stores/queryStore.ts
index 5ed72d52f..958f2248a 100644
--- a/src/stores/queryStore.ts
+++ b/src/stores/queryStore.ts
@@ -50,6 +50,23 @@ export const useQueryStore = defineStore("query", () => {
     if (tab) tab.sql = sql;
   }
 
+  function updateDatabase(id: string, database: string) {
+    const tab = tabs.value.find((t) => t.id === id);
+    if (!tab || tab.database === database) return;
+    tab.database = database;
+    tab.result = undefined;
+    tab.tableMeta = undefined;
+  }
+
+  function updateConnection(id: string, connectionId: string, database = "") {
+    const tab = tabs.value.find((t) => t.id === id);
+    if (!tab || tab.connectionId === connectionId) return;
+    tab.connectionId = connectionId;
+    tab.database = database;
+    tab.result = undefined;
+    tab.tableMeta = undefined;
+  }
+
   function setTableMeta(id: string, meta: NonNullable<QueryTab["tableMeta"]>) {
     const tab = tabs.value.find((t) => t.id === id);
     if (tab) tab.tableMeta = meta;
@@ -89,6 +106,8 @@ export const useQueryStore = defineStore("query", () => {
     createTab,
     closeTab,
     updateSql,
+    updateDatabase,
+    updateConnection,
     setTableMeta,
     executeCurrentTab,
   };
diff --git a/src/stores/settingsStore.ts b/src/stores/settingsStore.ts
index c4a0f680a..521420ec7 100644
--- a/src/stores/settingsStore.ts
+++ b/src/stores/settingsStore.ts
@@ -1,5 +1,6 @@
 import { defineStore } from "pinia";
 import { ref } from "vue";
+import * as api from "@/lib/tauri";
 
 export type AiProvider = "claude" | "openai" | "custom";
 
@@ -17,22 +18,35 @@ const defaultConfigs: Record<AiProvider, Omit<AiConfig, "apiKey">> = {
 };
 
 export const useSettingsStore = defineStore("settings", () => {
-  const saved = localStorage.getItem("dbx-ai-config");
-  const aiConfig = ref<AiConfig>(saved ? JSON.parse(saved) : { ...defaultConfigs.claude, apiKey: "" });
+  const aiConfig = ref<AiConfig>({ ...defaultConfigs.claude, apiKey: "" });
+  const isAiConfigLoaded = ref(false);
+
+  async function initAiConfig() {
+    if (isAiConfigLoaded.value) return;
+    const legacy = localStorage.getItem("dbx-ai-config");
+    const saved = await api.loadAiConfig().catch(() => null);
+    if (saved) {
+      aiConfig.value = saved;
+    } else if (legacy) {
+      aiConfig.value = JSON.parse(legacy);
+      await api.saveAiConfig(aiConfig.value).catch(() => {});
+      localStorage.removeItem("dbx-ai-config");
+    }
+    isAiConfigLoaded.value = true;
+  }
 
   function updateAiConfig(config: Partial<AiConfig>) {
-    Object.assign(aiConfig.value, config);
-    if (config.provider && config.provider !== aiConfig.value.provider) {
-      const defaults = defaultConfigs[config.provider];
-      aiConfig.value.endpoint = defaults.endpoint;
-      aiConfig.value.model = defaults.model;
+    const previousProvider = aiConfig.value.provider;
+    if (config.provider && config.provider !== previousProvider) {
+      Object.assign(aiConfig.value, defaultConfigs[config.provider]);
     }
-    localStorage.setItem("dbx-ai-config", JSON.stringify(aiConfig.value));
+    Object.assign(aiConfig.value, config);
+    api.saveAiConfig(aiConfig.value).catch(() => {});
   }
 
   function isConfigured(): boolean {
     return !!aiConfig.value.apiKey && !!aiConfig.value.endpoint;
   }
 
-  return { aiConfig, updateAiConfig, isConfigured };
+  return { aiConfig, isAiConfigLoaded, initAiConfig, updateAiConfig, isConfigured };
 });