feat(ai): agent loop Phase 2 — SQL execution tools, risk classification, and streaming tool calls

feat(ai): agent loop Phase 2 — SQL execution tools, risk classification, and streaming tool calls

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');

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
+}

crates/dbx-core/src/db/postgres.rs

@@ -1093,6 +1093,7 @@ pub async fn list_databases(pool: &Pool) -> Result<Vec<DatabaseInfo>, String> {
 }
 
 pub async fn list_tables(pool: &Pool, schema: &str) -> Result<Vec<TableInfo>, String> {
+    let schema = if schema.is_empty() { "public" } else { schema };
     let client = pool.get().await.map_err(|e| e.to_string())?;
     let stmt = client.prepare_cached(postgres_tables_sql()).await.map_err(|e| e.to_string())?;
     let rows = client.query(&stmt, &[&schema]).await.map_err(|e| e.to_string())?;
@@ -1356,6 +1357,7 @@ async fn get_columns_with_sql(
 }
 
 pub async fn get_columns(pool: &Pool, schema: &str, table: &str) -> Result<Vec<ColumnInfo>, String> {
+    let schema = if schema.is_empty() { "public" } else { schema };
     let client = pool.get().await.map_err(|e| e.to_string())?;
     match get_columns_with_sql(&client, POSTGRES_COLUMNS_SQL, schema, table).await {
         Ok(columns) => Ok(columns),

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;

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);
+    }
+}