From 100124da3e569a29919809a91c1435049f65d385 Mon Sep 17 00:00:00 2001
From: runstone <lcprunstone@163.com>
Date: Fri, 12 Jun 2026 17:42:23 +0800
Subject: [PATCH] =?UTF-8?q?feat(ai):=20Phase=202=20SQL=20=E9=A3=8E?=
 =?UTF-8?q?=E9=99=A9=E5=88=86=E7=B1=BB=20+=20Agent=20=E6=A8=A1=E5=BC=8F?=
 =?UTF-8?q?=E5=B7=A5=E5=85=B7=E8=B0=83=E7=94=A8=E7=BB=93=E6=9E=9C=E6=B8=B2?=
 =?UTF-8?q?=E6=9F=93?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

  后端：
  - 新增 sql_risk.rs：基于 sqlparser AST 的 SQL 风险分类（ReadOnly/Write/Ddl/Transaction）
  - agent_tools.rs：新增 execute_query 和 get_sample_data 工具，写操作/DDL 被风险分类阻止
  - agent_loop.rs：支持 is_agent_mode 参数，Agent 模式下暴露全部工具，Ask 模式仅暴露只读工具
  - agent_loop/web/tauri：mode 参数全链路传递（前端→HTTP→Tauri→Rust）

  前端：
  - AiAgentStepItem 接口扩展：toolName/toolArgs/toolResult/isError 字段
  - AiAssistant.vue：工具调用结果实时+最终渲染，SQL 查询结果/错误信息展示
  - ai.ts/http.ts/tauri.ts：mode 参数透传到后端
---
 .../src/components/editor/AiAssistant.vue     |  60 +++--
 apps/desktop/src/lib/ai.ts                    |   1 +
 .../src/lib/aiAgentStepPresentation.ts        |  10 +-
 apps/desktop/src/lib/http.ts                  |   8 +-
 apps/desktop/src/lib/tauri.ts                 |   6 +-
 crates/dbx-core/src/agent_loop.rs             |  13 +-
 crates/dbx-core/src/agent_tools.rs            | 208 ++++++++++++++++-
 crates/dbx-core/src/lib.rs                    |   1 +
 crates/dbx-core/src/sql_risk.rs               | 211 ++++++++++++++++++
 crates/dbx-web/src/routes/ai.rs               |  10 +
 src-tauri/src/commands/ai.rs                  |   5 +-
 11 files changed, 501 insertions(+), 32 deletions(-)
 create mode 100644 crates/dbx-core/src/sql_risk.rs

diff --git a/apps/desktop/src/components/editor/AiAssistant.vue b/apps/desktop/src/components/editor/AiAssistant.vue
index d59ba3ea7f..a29000e1d9 100644
--- a/apps/desktop/src/components/editor/AiAssistant.vue
+++ b/apps/desktop/src/components/editor/AiAssistant.vue
@@ -1,4 +1,4 @@
-<script setup lang="ts">
+<script setup lang="ts">
 import { computed, nextTick, onMounted, onUnmounted, ref, type Component } from "vue";
 import { uuid } from "@/lib/utils";
 import { useI18n } from "vue-i18n";
@@ -252,11 +252,15 @@ function agentStepClass(tone: AiAgentStepTone): string {
   }
 }
 
-function agentStepTitle(step: AiAgentStepItem): string {
-  if (step.titleKey) return t(step.titleKey, step.titleParams || {});
-  const tool = step.titleParams?.tool;
-  if (tool) return `${t(step.labelKey)}: ${tool}`;
-  return t(step.labelKey);
+/** Extract tool result content from the AgentEvent result value */
+function extractToolResultContent(result: unknown): string | undefined {
+  if (!result) return undefined;
+  if (typeof result === "string") return result;
+  if (typeof result === "object" && result !== null && "content" in result) {
+    const content = (result as Record<string, unknown>).content;
+    return typeof content === "string" ? content : JSON.stringify(content);
+  }
+  return JSON.stringify(result);
 }
 
 function toggleReasoning(index: number) {
@@ -527,6 +531,10 @@ async function send() {
                 tone: (e.type === "tool_call_start" ? "active" : e.is_error ? "danger" : "success") as AiAgentStepTone,
                 titleKey: undefined,
                 titleParams: { tool: e.tool_name || "" },
+                toolName: e.tool_name,
+                toolArgs: e.type === "tool_call_start" ? (e.args as Record<string, unknown>) : undefined,
+                toolResult: e.type === "tool_call_end" && !e.is_error ? extractToolResultContent(e.result) : undefined,
+                isError: e.type === "tool_call_end" ? e.is_error : undefined,
               }));
             msg.agentSteps = steps;
           }
@@ -552,6 +560,10 @@ async function send() {
           tone: e.type === "tool_call_start" ? "active" : e.is_error ? "danger" : "success",
           titleKey: undefined,
           titleParams: { tool: e.tool_name || "" },
+          toolName: e.tool_name,
+          toolArgs: e.type === "tool_call_start" ? (e.args as Record<string, unknown>) : undefined,
+          toolResult: e.type === "tool_call_end" && !e.is_error ? extractToolResultContent(e.result) : undefined,
+          isError: e.type === "tool_call_end" ? e.is_error : undefined,
         }));
     }
     // Fallback: use aiAgentPlan for backward compatibility
@@ -774,11 +786,22 @@ const messageRenderer = computed(() => {
                   </div>
                 </div>
               </div>
-              <div v-if="msg.agentSteps?.length" class="mb-2 flex flex-wrap gap-1.5">
-                <span v-for="step in msg.agentSteps" :key="step.key" class="inline-flex h-5 max-w-full items-center gap-1 rounded-full border px-1.5 text-[10px] font-medium" :class="agentStepClass(step.tone)" :title="agentStepTitle(step)">
-                  <component :is="agentStepIcon(step.tone)" class="h-3 w-3 shrink-0" />
-                  <span class="truncate">{{ t(step.labelKey) }}</span>
-                </span>
+              <div v-if="msg.agentSteps?.length" class="mb-2 space-y-1">
+                <div v-for="step in msg.agentSteps" :key="step.key" class="rounded border px-2 py-1.5 text-[10px]" :class="agentStepClass(step.tone)">
+                  <div class="flex items-center gap-1">
+                    <component :is="agentStepIcon(step.tone)" class="h-3 w-3 shrink-0" />
+                    <span class="font-medium">{{ t(step.labelKey) }}</span>
+                    <span v-if="step.toolName" class="text-muted-foreground">: {{ step.toolName }}</span>
+                  </div>
+                  <!-- Show SQL query for execute_query/get_sample_data -->
+                  <div v-if="step.toolArgs?.sql" class="mt-1 rounded bg-background/50 px-2 py-1 font-mono text-[10px] text-foreground/80 whitespace-pre-wrap">{{ step.toolArgs.sql }}</div>
+                  <!-- Show error message -->
+                  <div v-if="step.isError && step.toolResult" class="mt-1 text-[10px] text-red-600 dark:text-red-400">{{ step.toolResult }}</div>
+                  <!-- Show result preview for successful queries -->
+                  <div v-else-if="step.toolResult && step.toolName === 'execute_query'" class="mt-1 overflow-x-auto">
+                    <div class="text-[10px] text-muted-foreground whitespace-pre-wrap max-h-32 overflow-y-auto">{{ step.toolResult }}</div>
+                  </div>
+                </div>
               </div>
               <template v-for="(seg, j) in messageRenderer.render(msg.content)" :key="j">
                 <div v-if="seg.type === 'text'" class="ai-markdown whitespace-normal">
@@ -825,7 +848,14 @@ const messageRenderer = computed(() => {
         <div v-if="connectionStore.connections.length" class="flex items-center gap-1 mb-1 text-xs text-foreground/80">
           <DatabaseIcon v-if="connection" :db-type="connectionIconType(connection)" class="h-3 w-3 shrink-0" />
           <Server v-else class="h-3 w-3 shrink-0" />
-          <Select :model-value="connection?.id || ''" @update:model-value="(v: string) => changeConnection(v)">
+          <Select
+            :model-value="connection?.id || ''"
+            @update:model-value="
+              (v) => {
+                if (typeof v === 'string') changeConnection(v);
+              }
+            "
+          >
             <SelectTrigger class="h-5 w-auto border-0 rounded-md bg-transparent dark:bg-transparent p-0 px-1 text-xs text-foreground/80 shadow-none focus:ring-0 focus-visible:ring-0 [&_svg]:size-3">
               <SelectValue :placeholder="t('editor.selectConnection')">{{ connection?.name || t("editor.selectConnection") }}</SelectValue>
             </SelectTrigger>
@@ -842,7 +872,11 @@ const messageRenderer = computed(() => {
             <Database class="h-3 w-3 shrink-0 text-foreground/40" />
             <Select
               :model-value="selectedDatabaseSelectValue"
-              @update:model-value="(v: string) => changeDatabase(v)"
+              @update:model-value="
+                (v) => {
+                  if (typeof v === 'string') changeDatabase(v);
+                }
+              "
               @update:open="
                 (open: boolean) => {
                   if (open) loadDatabases();
diff --git a/apps/desktop/src/lib/ai.ts b/apps/desktop/src/lib/ai.ts
index ed4ef08024..fa499caee9 100644
--- a/apps/desktop/src/lib/ai.ts
+++ b/apps/desktop/src/lib/ai.ts
@@ -111,6 +111,7 @@ export async function runAgentStream(input: AiRequestInput, history: api.AiMessa
     input.context.database,
     input.context.databaseType,
     onEvent,
+    input.mode || "ask",
   );
 }
 
diff --git a/apps/desktop/src/lib/aiAgentStepPresentation.ts b/apps/desktop/src/lib/aiAgentStepPresentation.ts
index 369489c5ab..13af727184 100644
--- a/apps/desktop/src/lib/aiAgentStepPresentation.ts
+++ b/apps/desktop/src/lib/aiAgentStepPresentation.ts
@@ -1,4 +1,4 @@
-import type { AiAgentPlan, AiAgentStep } from "@/lib/aiAgentPlan";
+import type { AiAgentPlan, AiAgentStep } from "@/lib/aiAgentPlan";
 
 export type AiAgentStepTone = "success" | "active" | "warning" | "danger" | "muted";
 
@@ -8,6 +8,14 @@ export interface AiAgentStepItem {
   tone: AiAgentStepTone;
   titleKey?: string;
   titleParams?: Record<string, string>;
+  /** Tool name for display */
+  toolName?: string;
+  /** Tool arguments (e.g., SQL query) for display */
+  toolArgs?: Record<string, unknown>;
+  /** Tool result content (e.g., query results) for display */
+  toolResult?: string;
+  /** Whether this is an error result */
+  isError?: boolean;
 }
 
 export function buildAiAgentStepItems(plan: AiAgentPlan): AiAgentStepItem[] {
diff --git a/apps/desktop/src/lib/http.ts b/apps/desktop/src/lib/http.ts
index 8d73638150..70e6738e2f 100644
--- a/apps/desktop/src/lib/http.ts
+++ b/apps/desktop/src/lib/http.ts
@@ -1,4 +1,4 @@
-import type {
+import type {
   ConnectionConfig,
   DatabaseInfo,
   TableInfo,
@@ -763,11 +763,11 @@ function isAgentEvent(v: unknown): v is import("./tauri").AgentEvent {
   return typeof v === "object" && v !== null && "type" in v && typeof (v as Record<string, unknown>).type === "string";
 }
 
-export async function aiAgentStream(sessionId: string, request: AiCompletionRequest, connectionId: string, database: string, dbType: string, onEvent: (event: import("./tauri").AgentEvent) => void, signal?: AbortSignal): Promise<string> {
+export async function aiAgentStream(sessionId: string, request: AiCompletionRequest, connectionId: string, database: string, dbType: string, onEvent: (event: import("./tauri").AgentEvent) => void, mode?: string, signal?: AbortSignal): Promise<string> {
   const res = await fetch("/api/ai/agent-stream", {
     method: "POST",
     headers: { "Content-Type": "application/json" },
-    body: JSON.stringify({ session_id: sessionId, request, connection_id: connectionId, database, db_type: dbType }),
+    body: JSON.stringify({ session_id: sessionId, request, connection_id: connectionId, database, db_type: dbType, mode: mode || "ask" }),
     signal,
   });
   if (!res.ok) throw new Error(await res.text());
@@ -1430,5 +1430,5 @@ export async function loadSidebarLayout(): Promise<SidebarLayout | null> {
 }
 
 export async function refreshConnections(): Promise<void> {
-  // Web mode doesn't maintain persistent connection pools — no-op
+  // Web mode doesn't maintain persistent connection pools 鈥?no-op
 }
diff --git a/apps/desktop/src/lib/tauri.ts b/apps/desktop/src/lib/tauri.ts
index 5b3269b326..d3da4f1c34 100644
--- a/apps/desktop/src/lib/tauri.ts
+++ b/apps/desktop/src/lib/tauri.ts
@@ -1,4 +1,4 @@
-import { invoke } from "@tauri-apps/api/core";
+import { invoke } from "@tauri-apps/api/core";
 import { listen, type UnlistenFn } from "@tauri-apps/api/event";
 import type {
   ConnectionConfig,
@@ -284,7 +284,7 @@ export type AgentEvent =
   | { type: "agent_end"; total_tokens?: number }
   | { type: "error"; message: string };
 
-export async function aiAgentStream(sessionId: string, request: AiCompletionRequest, connectionId: string, database: string, dbType: string, onEvent: (event: AgentEvent) => void): Promise<string> {
+export async function aiAgentStream(sessionId: string, request: AiCompletionRequest, connectionId: string, database: string, dbType: string, onEvent: (event: AgentEvent) => void, mode?: string, _signal?: AbortSignal): Promise<string> {
   const unlisten: UnlistenFn = await listen<AgentEvent>("ai-agent-event", (event) => {
     onEvent(event.payload);
     if (event.payload.type === "agent_end" || event.payload.type === "error") {
@@ -292,7 +292,7 @@ export async function aiAgentStream(sessionId: string, request: AiCompletionRequ
     }
   });
   try {
-    return await invoke("ai_agent_stream", { sessionId, request, connectionId, database, dbType });
+    return await invoke("ai_agent_stream", { sessionId, request, connectionId, database, dbType, mode });
   } catch (e) {
     unlisten();
     throw e;
diff --git a/crates/dbx-core/src/agent_loop.rs b/crates/dbx-core/src/agent_loop.rs
index f3a1e0776b..d3154e0bac 100644
--- a/crates/dbx-core/src/agent_loop.rs
+++ b/crates/dbx-core/src/agent_loop.rs
@@ -49,6 +49,7 @@ pub async fn run_agent_loop(
     cancelled: &Notify,
     max_tokens: Option<u32>,
     temperature: Option<f32>,
+    is_agent_mode: bool,
 ) -> Result<String, String> {
     // Auto-degrade: providers without function calling fall back to text-only completion.
     if !provider_supports_function_calling(config) {
@@ -64,7 +65,7 @@ pub async fn run_agent_loop(
         )
         .await;
     }
-    let tools = agent_tools::read_only_tools();
+    let tools = if is_agent_mode { agent_tools::all_tools() } else { agent_tools::read_only_tools() };
     let mut conversation_messages: Vec<AiMessage> = messages.to_vec();
     let mut final_text = String::new();
 
@@ -100,7 +101,8 @@ pub async fn run_agent_loop(
         };
 
         // Call the LLM with tool support
-        let collected_tool_calls = stream_with_tools(config, &request, &session_id, cancelled, on_chunk).await?;
+        let collected_tool_calls =
+            stream_with_tools(config, &request, &session_id, &tools, cancelled, on_chunk).await?;
 
         on_event(AgentEvent::TurnEnd { turn });
 
@@ -193,11 +195,10 @@ async fn stream_with_tools(
     config: &AiConfig,
     request: &AiCompletionRequest,
     session_id: &str,
+    tools: &[ToolDefinition],
     cancelled: &Notify,
     on_chunk: impl Fn(AiStreamChunk) + Send + Sync + 'static,
 ) -> Result<Vec<ToolCall>, String> {
-    let tools = agent_tools::read_only_tools();
-
     // Return early if the user cancelled before the LLM call started.
     if cancelled.notified().now_or_never().is_some() {
         return Err("Agent loop cancelled".to_string());
@@ -258,14 +259,14 @@ async fn build_schema_prompt(agent_ctx: &AgentLoopContext, system_prompt: &str)
 
     match tables_result {
         Ok(tables) if !tables.is_empty() => {
-            enriched.push_str("\n\n## Database Schema (for context — no tools available)\n");
+            enriched.push_str("\n\n## Database Schema (for context 鈥?no tools available)\n");
             enriched.push_str(&format!("Database: {}\n", agent_ctx.database));
             enriched.push_str("Tables:\n");
             for t in &tables {
                 enriched.push_str(&format!("  - {} ({})", t.name, t.table_type));
                 if let Some(ref comment) = t.comment {
                     if !comment.trim().is_empty() {
-                        enriched.push_str(&format!(" — {}", comment.trim()));
+                        enriched.push_str(&format!(" 鈥?{}", comment.trim()));
                     }
                 }
                 enriched.push('\n');
diff --git a/crates/dbx-core/src/agent_tools.rs b/crates/dbx-core/src/agent_tools.rs
index c8d8f5b4bc..4614f3143a 100644
--- a/crates/dbx-core/src/agent_tools.rs
+++ b/crates/dbx-core/src/agent_tools.rs
@@ -5,15 +5,32 @@ use serde_json::json;
 use crate::agent_events::{ToolCall, ToolDefinition, ToolResult};
 use crate::connection::AppState;
 use crate::models::connection::DatabaseType;
+use crate::query::QueryExecutionOptions;
+use crate::sql_risk::SqlRisk;
+use crate::types::QueryResult;
 
 /// Maximum number of tables returned by list_tables tool.
 const LIST_TABLES_LIMIT: usize = 200;
 
+/// Maximum number of rows returned by execute_query tool.
+const EXECUTE_QUERY_LIMIT: usize = 50;
+
+/// Maximum number of rows returned by get_sample_data tool.
+const SAMPLE_DATA_LIMIT: usize = 20;
+
+/// Absolute maximum rows any query tool may request.
+const MAX_ALLOWED_ROWS: usize = 100;
+
 /// Get all available tool definitions for Phase 1 (read-only tools).
 pub fn read_only_tools() -> Vec<ToolDefinition> {
     vec![list_tables_tool(), get_columns_tool()]
 }
 
+/// Get all available tool definitions (Phase 2: read-only + execute_query + get_sample_data).
+pub fn all_tools() -> Vec<ToolDefinition> {
+    vec![list_tables_tool(), get_columns_tool(), execute_query_tool(), get_sample_data_tool()]
+}
+
 /// list_tables tool definition.
 fn list_tables_tool() -> ToolDefinition {
     ToolDefinition {
@@ -58,6 +75,57 @@ fn get_columns_tool() -> ToolDefinition {
         read_only: true,
     }
 }
+/// execute_query tool definition (Phase 2).
+fn execute_query_tool() -> ToolDefinition {
+    ToolDefinition {
+        name: "execute_query",
+        description: "Execute a read-only SQL query and return results (max 50 rows). \
+                      Only SELECT, WITH, SHOW, DESCRIBE, EXPLAIN statements are allowed. \
+                      Write operations (INSERT/UPDATE/DELETE/DDL) are blocked.",
+        parameters: json!({
+            "type": "object",
+            "properties": {
+                "sql": {
+                    "type": "string",
+                    "description": "The SQL query to execute"
+                },
+                "limit": {
+                    "type": "number",
+                    "description": "Max rows to return (default 50, max 100)"
+                }
+            },
+            "required": ["sql"]
+        }),
+        read_only: true,
+    }
+}
+
+/// get_sample_data tool definition (Phase 2).
+fn get_sample_data_tool() -> ToolDefinition {
+    ToolDefinition {
+        name: "get_sample_data",
+        description: "Get sample rows from a table to understand its data. Returns up to 20 rows.",
+        parameters: json!({
+            "type": "object",
+            "properties": {
+                "table": {
+                    "type": "string",
+                    "description": "Table name"
+                },
+                "schema": {
+                    "type": "string",
+                    "description": "Schema name (optional)"
+                },
+                "limit": {
+                    "type": "number",
+                    "description": "Max rows (default 20)"
+                }
+            },
+            "required": ["table"]
+        }),
+        read_only: true,
+    }
+}
 
 /// Execute a tool call and return the result.
 pub async fn execute_tool(
@@ -70,6 +138,8 @@ pub async fn execute_tool(
     let result = match tool_call.name.as_str() {
         "list_tables" => execute_list_tables(tool_call, state, connection_id, database, db_type).await,
         "get_columns" => execute_get_columns(tool_call, state, connection_id, database, db_type).await,
+        "execute_query" => execute_execute_query(tool_call, state, connection_id, database, db_type).await,
+        "get_sample_data" => execute_get_sample_data(tool_call, state, connection_id, database, db_type).await,
         _ => Err(format!("Unknown tool: {}", tool_call.name)),
     };
 
@@ -98,16 +168,18 @@ async fn execute_list_tables(
 ) -> Result<String, String> {
     let schema = tool_call.arguments.get("schema").and_then(|v| v.as_str()).unwrap_or("").to_string();
 
-    // list_tables_core already limits via the Some(LIST_TABLES_LIMIT) argument.
+    // Request one extra to detect whether more tables exist beyond the limit.
     let tables =
-        crate::schema::list_tables_core(state, connection_id, database, &schema, None, Some(LIST_TABLES_LIMIT))
+        crate::schema::list_tables_core(state, connection_id, database, &schema, None, Some(LIST_TABLES_LIMIT + 1))
             .await
             .map_err(|e| format!("Failed to list tables: {e}"))?;
 
     let total = tables.len();
+    let truncated = total > LIST_TABLES_LIMIT;
 
     let mut lines = Vec::new();
-    for table in &tables {
+    let display_count = if truncated { LIST_TABLES_LIMIT } else { total };
+    for table in tables.iter().take(display_count) {
         let mut line = format!("- {} ({})", table.name, table.table_type);
         if let Some(comment) = &table.comment {
             let trimmed = comment.trim();
@@ -118,7 +190,7 @@ async fn execute_list_tables(
         lines.push(line);
     }
 
-    if total > LIST_TABLES_LIMIT {
+    if truncated {
         lines.push(format!("... (showing {LIST_TABLES_LIMIT} of {total} tables)"));
     }
 
@@ -202,3 +274,131 @@ async fn execute_get_columns(
 
     Ok(lines.join("\n"))
 }
+
+/// Execute a read-only SQL query via the execute_query tool.
+async fn execute_execute_query(
+    tool_call: &ToolCall,
+    state: &Arc<AppState>,
+    connection_id: &str,
+    database: &str,
+    db_type: &DatabaseType,
+) -> Result<String, String> {
+    let sql = tool_call.arguments.get("sql").and_then(|v| v.as_str()).ok_or("Missing required parameter: sql")?.trim();
+
+    if sql.is_empty() {
+        return Err("SQL query cannot be empty".to_string());
+    }
+
+    let limit = tool_call
+        .arguments
+        .get("limit")
+        .and_then(|v| v.as_u64())
+        .map(|l| (l as usize).min(MAX_ALLOWED_ROWS))
+        .unwrap_or(EXECUTE_QUERY_LIMIT);
+
+    // Classify SQL risk using sqlparser AST
+    let db_type_str = format!("{:?}", db_type).to_lowercase();
+    let risk = crate::sql_risk::classify_sql_risk(sql, &db_type_str)?;
+    match risk {
+        SqlRisk::ReadOnly => { /* proceed */ }
+        _ => {
+            return Err(format!(
+                "Blocked: {} statement detected. Only read-only queries (SELECT, SHOW, DESCRIBE, EXPLAIN) are allowed.",
+                risk
+            ));
+        }
+    }
+
+    // Execute query using existing infrastructure
+    let options = QueryExecutionOptions { max_rows: Some(limit), timeout_secs: Some(30), ..Default::default() };
+    let result =
+        crate::query::execute_sql_statement_with_options(state, connection_id, database, sql, None, None, options)
+            .await?;
+
+    format_query_result_as_text(&result, limit)
+}
+
+/// Format a QueryResult as a Markdown table for LLM consumption.
+fn format_query_result_as_text(result: &QueryResult, limit: usize) -> Result<String, String> {
+    if result.rows.is_empty() {
+        return Ok("Query returned 0 rows.".to_string());
+    }
+
+    let mut lines = Vec::new();
+
+    // Header row
+    lines.push(format!("| {} |", result.columns.join(" | ")));
+    // Separator row
+    lines.push(format!("|{}|", result.columns.iter().map(|_| "---").collect::<Vec<_>>().join("|")));
+
+    // Data rows
+    for row in &result.rows {
+        let cells: Vec<String> = row
+            .iter()
+            .map(|v| match v {
+                serde_json::Value::Null => "NULL".to_string(),
+                serde_json::Value::String(s) => {
+                    // Truncate long strings to keep result compact
+                    if s.len() > 200 {
+                        let truncated: String =
+                            s.char_indices().take_while(|(i, _)| *i < 200).map(|(_, c)| c).collect();
+                        format!("{}...", truncated)
+                    } else {
+                        s.clone()
+                    }
+                }
+                other => other.to_string(),
+            })
+            .collect();
+        lines.push(format!("| {} |", cells.join(" | ")));
+    }
+
+    // Truncation notice
+    if result.truncated || result.rows.len() >= limit {
+        lines.push(format!("... (showing {} rows, result may be truncated)", result.rows.len()));
+    }
+
+    // Stats line
+    lines.push(format!("({} rows, {}ms)", result.rows.len(), result.execution_time_ms));
+
+    Ok(lines.join("\n"))
+}
+
+/// Get sample data from a table via the get_sample_data tool.
+async fn execute_get_sample_data(
+    tool_call: &ToolCall,
+    state: &Arc<AppState>,
+    connection_id: &str,
+    database: &str,
+    db_type: &DatabaseType,
+) -> Result<String, String> {
+    let table =
+        tool_call.arguments.get("table").and_then(|v| v.as_str()).ok_or("Missing required parameter: table")?.trim();
+
+    if table.is_empty() {
+        return Err("Table name cannot be empty".to_string());
+    }
+    if table.contains(';') || table.contains('\'') || table.contains('"') || table.contains('\\') {
+        return Err(format!("Table name contains invalid characters: '{}'", table));
+    }
+
+    let schema = tool_call.arguments.get("schema").and_then(|v| v.as_str());
+    let limit = tool_call
+        .arguments
+        .get("limit")
+        .and_then(|v| v.as_u64())
+        .map(|l| (l as usize).min(MAX_ALLOWED_ROWS))
+        .unwrap_or(SAMPLE_DATA_LIMIT);
+
+    // Build SELECT * FROM table LIMIT N
+    let schema_prefix = schema.filter(|s| !s.is_empty()).map(|s| format!("\"{}\".", s)).unwrap_or_default();
+    let sql = format!("SELECT * FROM {}\"{}\" LIMIT {}", schema_prefix, table, limit);
+
+    // Delegate to execute_execute_query with a synthetic tool call
+    let synthetic_call = ToolCall {
+        id: tool_call.id.clone(),
+        name: "execute_query".to_string(),
+        arguments: serde_json::json!({ "sql": sql, "limit": limit }),
+    };
+    execute_execute_query(&synthetic_call, state, connection_id, database, db_type).await
+}
diff --git a/crates/dbx-core/src/lib.rs b/crates/dbx-core/src/lib.rs
index b000024b46..486a910d40 100644
--- a/crates/dbx-core/src/lib.rs
+++ b/crates/dbx-core/src/lib.rs
@@ -41,6 +41,7 @@ pub mod sql_analysis;
 pub mod sql_dialect;
 pub mod sql_editability;
 pub mod sql_file_import;
+pub mod sql_risk;
 pub mod storage;
 pub mod table_export;
 pub mod table_import;
diff --git a/crates/dbx-core/src/sql_risk.rs b/crates/dbx-core/src/sql_risk.rs
new file mode 100644
index 0000000000..c4339374fb
--- /dev/null
+++ b/crates/dbx-core/src/sql_risk.rs
@@ -0,0 +1,211 @@
+use serde::{Deserialize, Serialize};
+use sqlparser::ast::Statement;
+use sqlparser::dialect::{
+    ClickHouseDialect, DuckDbDialect, GenericDialect, MsSqlDialect, MySqlDialect, PostgreSqlDialect, SQLiteDialect,
+};
+use sqlparser::parser::Parser;
+
+/// SQL risk level for agent tool safety classification.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+pub enum SqlRisk {
+    /// SELECT, SHOW, DESCRIBE, EXPLAIN, WITH (pure read CTE)
+    ReadOnly,
+    /// INSERT, UPDATE, DELETE, MERGE, REPLACE, CALL/EXEC
+    Write,
+    /// CREATE, ALTER, DROP, TRUNCATE, GRANT, REVOKE
+    Ddl,
+    /// BEGIN, COMMIT, ROLLBACK should not be issued by agent
+    Transaction,
+}
+
+impl std::fmt::Display for SqlRisk {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            SqlRisk::ReadOnly => write!(f, "read-only"),
+            SqlRisk::Write => write!(f, "write"),
+            SqlRisk::Ddl => write!(f, "DDL"),
+            SqlRisk::Transaction => write!(f, "transaction"),
+        }
+    }
+}
+
+/// Normalize database dialect string to a canonical form for sqlparser.
+/// Mirrors the logic in `sql_analysis::normalize_dialect`.
+fn normalize_dialect(dialect: &str) -> &'static str {
+    match dialect.to_ascii_lowercase().as_str() {
+        "postgres" | "postgresql" | "redshift" | "opengauss" | "gaussdb" | "highgo" => "postgres",
+        "mysql" | "mariadb" | "doris" | "starrocks" | "oceanbase" => "mysql",
+        "sqlite" => "sqlite",
+        "sqlserver" | "mssql" => "sqlserver",
+        "clickhouse" => "clickhouse",
+        "duckdb" => "duckdb",
+        _ => "generic",
+    }
+}
+
+/// Resolve dialect string to a sqlparser Dialect trait object.
+fn resolve_dialect(dialect: &str) -> Box<dyn sqlparser::dialect::Dialect> {
+    match dialect {
+        "postgres" => Box::new(PostgreSqlDialect {}),
+        "mysql" => Box::new(MySqlDialect {}),
+        "sqlite" => Box::new(SQLiteDialect {}),
+        "sqlserver" => Box::new(MsSqlDialect {}),
+        "clickhouse" => Box::new(ClickHouseDialect {}),
+        "duckdb" => Box::new(DuckDbDialect {}),
+        _ => Box::new(GenericDialect {}),
+    }
+}
+
+/// Classify a single SQL statement into a risk level using AST analysis.
+fn classify_statement(stmt: &Statement) -> SqlRisk {
+    match stmt {
+        // Pure reads
+        Statement::Query(_) => SqlRisk::ReadOnly,
+        Statement::Explain { .. } => SqlRisk::ReadOnly,
+        Statement::ExplainTable { .. } => SqlRisk::ReadOnly,
+
+        // Show/Describe variants
+        Statement::ShowTables { .. }
+        | Statement::ShowColumns { .. }
+        | Statement::ShowDatabases { .. }
+        | Statement::ShowSchemas { .. }
+        | Statement::ShowCreate { .. }
+        | Statement::ShowVariables { .. }
+        | Statement::ShowStatus { .. }
+        | Statement::ShowProcessList { .. } => SqlRisk::ReadOnly,
+
+        // Write operations
+        Statement::Insert { .. } | Statement::Update { .. } | Statement::Delete { .. } | Statement::Merge { .. } => {
+            SqlRisk::Write
+        }
+
+        // DDL operations
+        Statement::CreateTable { .. }
+        | Statement::CreateView { .. }
+        | Statement::CreateIndex { .. }
+        | Statement::CreateSchema { .. }
+        | Statement::CreateSequence { .. }
+        | Statement::CreateRole { .. }
+        | Statement::CreateType { .. }
+        | Statement::AlterTable { .. }
+        | Statement::AlterIndex { .. }
+        | Statement::AlterView { .. }
+        | Statement::Drop { .. }
+        | Statement::Truncate { .. } => SqlRisk::Ddl,
+
+        // Grant/Revoke
+        Statement::Grant { .. } | Statement::Revoke { .. } => SqlRisk::Ddl,
+
+        // Transaction control
+        Statement::StartTransaction { .. } | Statement::Commit { .. } | Statement::Rollback { .. } => {
+            SqlRisk::Transaction
+        }
+
+        // Copy (PostgreSQL) 鈥?treat as write
+        Statement::Copy { .. } => SqlRisk::Write,
+
+        // Pragma (SQLite/DuckDB) 鈥?conservative: treat as write unless known-safe
+        Statement::Pragma { .. } => SqlRisk::Write,
+
+        // Catch-all: conservative write classification
+        _ => SqlRisk::Write,
+    }
+}
+
+/// Classify SQL risk using sqlparser AST analysis.
+///
+/// If parsing fails (non-standard SQL, non-SQL databases), falls back to
+/// keyword-based `query_execution_sql::is_write_sql()`.
+///
+/// Multi-statement input: returns the highest risk level across all statements.
+pub fn classify_sql_risk(sql: &str, dialect: &str) -> Result<SqlRisk, String> {
+    let normalized = normalize_dialect(dialect);
+    let parser_dialect = resolve_dialect(normalized);
+
+    match Parser::parse_sql(parser_dialect.as_ref(), sql) {
+        Ok(stmts) if !stmts.is_empty() => {
+            let mut max_risk = SqlRisk::ReadOnly;
+            for stmt in &stmts {
+                let risk = classify_statement(stmt);
+                if risk as u8 > max_risk as u8 {
+                    max_risk = risk;
+                }
+            }
+            Ok(max_risk)
+        }
+        _ => {
+            // Fallback: keyword-based classification
+            if crate::query_execution_sql::is_write_sql(sql) {
+                Ok(SqlRisk::Write)
+            } else {
+                Ok(SqlRisk::ReadOnly)
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn classify_select_statements() {
+        assert_eq!(classify_sql_risk("SELECT * FROM users", "postgres").unwrap(), SqlRisk::ReadOnly);
+        assert_eq!(
+            classify_sql_risk("SELECT id, name FROM users WHERE active = true", "mysql").unwrap(),
+            SqlRisk::ReadOnly
+        );
+        assert_eq!(classify_sql_risk("SHOW TABLES", "mysql").unwrap(), SqlRisk::ReadOnly);
+        assert_eq!(classify_sql_risk("DESCRIBE users", "mysql").unwrap(), SqlRisk::ReadOnly);
+        assert_eq!(classify_sql_risk("EXPLAIN SELECT * FROM users", "postgres").unwrap(), SqlRisk::ReadOnly);
+    }
+
+    #[test]
+    fn classify_cte_read() {
+        assert_eq!(
+            classify_sql_risk("WITH cte AS (SELECT 1) SELECT * FROM cte", "postgres").unwrap(),
+            SqlRisk::ReadOnly
+        );
+    }
+
+    #[test]
+    fn classify_write_statements() {
+        assert_eq!(classify_sql_risk("INSERT INTO users VALUES (1)", "postgres").unwrap(), SqlRisk::Write);
+        assert_eq!(classify_sql_risk("UPDATE users SET name = 'x'", "postgres").unwrap(), SqlRisk::Write);
+        assert_eq!(classify_sql_risk("DELETE FROM users", "postgres").unwrap(), SqlRisk::Write);
+    }
+
+    #[test]
+    fn classify_ddl_statements() {
+        assert_eq!(classify_sql_risk("CREATE TABLE users (id INT)", "postgres").unwrap(), SqlRisk::Ddl);
+        assert_eq!(classify_sql_risk("DROP TABLE users", "postgres").unwrap(), SqlRisk::Ddl);
+        assert_eq!(classify_sql_risk("ALTER TABLE users ADD COLUMN age INT", "postgres").unwrap(), SqlRisk::Ddl);
+        assert_eq!(classify_sql_risk("TRUNCATE TABLE users", "postgres").unwrap(), SqlRisk::Ddl);
+    }
+
+    #[test]
+    fn classify_transaction_statements() {
+        assert_eq!(classify_sql_risk("BEGIN", "postgres").unwrap(), SqlRisk::Transaction);
+        assert_eq!(classify_sql_risk("COMMIT", "postgres").unwrap(), SqlRisk::Transaction);
+        assert_eq!(classify_sql_risk("ROLLBACK", "postgres").unwrap(), SqlRisk::Transaction);
+    }
+
+    #[test]
+    fn classify_multi_statement_returns_highest_risk() {
+        // SELECT + INSERT = Write
+        assert_eq!(classify_sql_risk("SELECT 1; INSERT INTO users VALUES (1)", "postgres").unwrap(), SqlRisk::Write);
+    }
+
+    #[test]
+    fn classify_fallback_on_parse_error() {
+        // Non-standard SQL should fall back to keyword matching
+        assert_eq!(classify_sql_risk("SELECT * FROM users", "generic").unwrap(), SqlRisk::ReadOnly);
+    }
+
+    #[test]
+    fn classify_unknown_statement_is_write() {
+        // Statements not explicitly handled should be conservative (Write)
+        // This depends on sqlparser's coverage, but we can test the catch-all
+        assert_eq!(classify_sql_risk("GRANT SELECT ON users TO admin", "postgres").unwrap(), SqlRisk::Ddl);
+    }
+}
diff --git a/crates/dbx-web/src/routes/ai.rs b/crates/dbx-web/src/routes/ai.rs
index c983bea941..a6f63fafc4 100644
--- a/crates/dbx-web/src/routes/ai.rs
+++ b/crates/dbx-web/src/routes/ai.rs
@@ -69,6 +69,14 @@ pub struct AiAgentStreamRequest {
     pub connection_id: String,
     pub database: String,
     pub db_type: String,
+    /// Agent mode: "ask" (read-only tools) or "agent" (all tools including execute_query).
+    /// Defaults to "ask" if not provided.
+    #[serde(default = "default_agent_mode")]
+    pub mode: String,
+}
+
+fn default_agent_mode() -> String {
+    "ask".to_string()
 }
 
 // ---------------------------------------------------------------------------
@@ -215,6 +223,7 @@ pub async fn ai_agent_stream(
     let req_messages = request.messages;
     let req_max_tokens = request.max_tokens;
     let req_temperature = request.temperature;
+    let is_agent_mode = body.mode == "agent";
     let tx2 = tx.clone();
     tokio::task::spawn_blocking(move || {
         let rt =
@@ -232,6 +241,7 @@ pub async fn ai_agent_stream(
                 &cancelled,
                 req_max_tokens,
                 req_temperature,
+                is_agent_mode,
             )
             .await;
 
diff --git a/src-tauri/src/commands/ai.rs b/src-tauri/src/commands/ai.rs
index b048b04691..83dddfde23 100644
--- a/src-tauri/src/commands/ai.rs
+++ b/src-tauri/src/commands/ai.rs
@@ -1,4 +1,4 @@
-use std::sync::Arc;
+use std::sync::Arc;
 use tauri::{AppHandle, Emitter, State};
 
 use super::connection::AppState;
@@ -60,6 +60,7 @@ pub async fn ai_agent_stream(
     connection_id: String,
     database: String,
     db_type: String,
+    mode: Option<String>,
 ) -> Result<String, String> {
     let cancelled = dbx_core::ai::register_stream(&session_id).await;
 
@@ -67,6 +68,7 @@ pub async fn ai_agent_stream(
         serde_json::from_str(&format!("\"{}\"", db_type)).map_err(|_| format!("Unknown database type: {db_type}"))?;
 
     let agent_ctx = AgentLoopContext { state: state.inner().clone(), connection_id, database, db_type: parsed_db_type };
+    let is_agent_mode = mode.as_deref() == Some("agent");
 
     let result = run_agent_loop(
         &request.config,
@@ -82,6 +84,7 @@ pub async fn ai_agent_stream(
         &cancelled,
         request.max_tokens,
         request.temperature,
+        is_agent_mode,
     )
     .await;