feat: rescope to AI-first DB harness with multi-DB chat agent

Removes enterprise/DBA features and replaces the marginal AI bar with a
central chat agent that has progressive-discovery tools, cross-session
memory, saved-query reuse, and inline result actions. Adds ClickHouse
support alongside PostgreSQL/Greenplum.

Cleanup
- Drop ~10k LOC of advanced features: Docker, Snapshots, Validation,
  Index Advisor, Role/User Management, Data Generator, ERD, Lookup.
- Trim deps: drop @xyflow/react, dagre, @types/dagre; cut tokio features
  to rt-multi-thread/sync/time/net/macros.
- Remove unused TuskError variants and dead helpers (topological_sort,
  invalidate_schema_cache).

Multi-DB (PostgreSQL + ClickHouse)
- New src-tauri/src/db/ module: ChClient (HTTP-based, reuses reqwest),
  sql_guard (cross-flavor read-only whitelist with 8 tests).
- ConnectionConfig gains db_flavor and secure fields with serde defaults
  for backwards-compatible connections.json.
- All connection/query/schema/data commands dispatch by flavor; CH
  covers connect, execute_query, list_databases/schemas/tables/views/
  columns/completion_schema, paginated table fetch.
- Frontend: dbCapabilities matrix, ConnectionDialog engine selector
  with port auto-swap and HTTPS toggle, SqlEditor switches to
  StandardSQL dialect for CH, TableDataView surfaces CH connections as
  read-only.

AI-first chat agent
- New src/components/chat/ panel with composer, message rendering,
  collapsible tool-call/result blocks, top-level ErrorBoundary.
- Backend agent loop in commands/chat.rs with strict-JSON tool
  protocol. Nine tools: list_databases, list_tables, get_columns,
  switch_database, run_query, remember, save_query, find_queries, final.
  Forgiving parser accepts both flat and nested-input shapes.
- Compressed history: only the last 4 run_query results carry sample
  rows (≤10, cells truncated to 200 chars) into LLM context; older
  results marked omitted.
- System prompt uses lite OVERVIEW (DB list + active-DB tables only)
  instead of full DDL — schema details are loaded on demand via
  get_columns. CH OVERVIEW shows cross-DB tables since CH allows
  db.table queries.

Cross-session memory (F1)
- Per-connection markdown file at app_data_dir/memory/<connection_id>.md,
  16KB cap with oldest-block eviction. Agent appends via remember()
  tool; the file is injected into LEARNED NOTES section of every system
  prompt.
- New Memory sidebar tab with editable textarea, badge for note count,
  empty-state with template. Edits picked up on the next agent turn.

Saved-query reuse (F2)
- Tools save_query and find_queries scoped to current connection.
  save_query attaches a UUID + timestamp; find_queries returns top 10
  matches with SQL preview ≤500 chars.
- Storage shared with the sidebar Saved panel.

Inline result actions (F3)
- run_query result block in chat gets Open-full (90vw × 80vh modal with
  full ResultsTable, no row cap) and Export (reuses ExportDialog for
  CSV/JSON via existing exportCsv/exportJson commands).

Verification
- cargo check clean, zero warnings.
- cargo test --lib: 50 pass (20 chat parser + 4 memory + 8 sql_guard +
  6 clean_sql + 12 escape_ident).
- npx tsc --noEmit clean.
- npx vitest run: 20 pass.
This commit is contained in:
2026-05-06 19:30:44 +03:00
parent 652937f7f5
commit 4f7afc17f4
89 changed files with 4151 additions and 10756 deletions

View File

@@ -0,0 +1,558 @@
//! Chat agent tool handlers (chat v2).
//!
//! Each `*_tool` function returns a plain string formatted for direct injection
//! into the LLM tool-result history. They reuse the schema helpers in
//! `commands::ai` and `commands::schema` rather than re-implementing SQL.
use crate::commands::ai::{
fetch_column_comments, fetch_columns, fetch_enum_types, fetch_foreign_keys_raw,
fetch_table_comments, fetch_unique_constraints, format_table_block, ColumnInfo,
};
use crate::commands::connections::{load_connection_config, switch_database_core};
use crate::commands::saved_queries::{list_saved_queries_core, save_query_core};
use crate::commands::schema::{list_databases_core, list_tables_core};
use crate::error::{TuskError, TuskResult};
use crate::models::saved_queries::SavedQuery;
use crate::state::{AppState, CachedVec, DbFlavor};
use sqlx::{PgPool, Row};
use std::collections::{BTreeMap, HashMap};
use std::time::{Duration, Instant};
use tauri::AppHandle;
const TOOL_CACHE_TTL: Duration = Duration::from_secs(300);
const MAX_TABLES_PER_GET_COLUMNS: usize = 20;
const COLUMNS_TOOL_OUTPUT_CAP: usize = 15_000;
// ---------------------------------------------------------------------------
// list_databases
// ---------------------------------------------------------------------------
pub async fn list_databases_tool(state: &AppState, connection_id: &str) -> TuskResult<String> {
let dbs = list_databases_core(state, connection_id).await?;
let active = active_db_name(state, connection_id).await;
let mut out = format!("DATABASES ({}):", dbs.len());
for db in &dbs {
if Some(db) == active.as_ref() {
out.push_str(&format!("\n * {} (active)", db));
} else {
out.push_str(&format!("\n {}", db));
}
}
Ok(out)
}
// ---------------------------------------------------------------------------
// list_tables
// ---------------------------------------------------------------------------
pub async fn list_tables_tool(
app: &AppHandle,
state: &AppState,
connection_id: &str,
db: Option<&str>,
) -> TuskResult<String> {
let active = active_db_name(state, connection_id).await;
let target = db.map(|s| s.to_string()).or_else(|| active.clone());
let target_name = match target.as_deref() {
Some(n) => n.to_string(),
None => return Err(TuskError::Custom("No active database selected.".into())),
};
let same_as_active = active.as_deref() == Some(target_name.as_str());
let flavor = state.get_flavor(connection_id).await;
let table_names = match (flavor, same_as_active) {
(DbFlavor::ClickHouse, _) => list_tables_clickhouse(state, connection_id, &target_name).await?,
(_, true) => list_tables_active_pg(state, connection_id).await?,
(_, false) => list_tables_other_pg(app, state, connection_id, &target_name).await?,
};
let header = if same_as_active {
format!("TABLES IN ACTIVE DATABASE `{}` ({}):", target_name, table_names.len())
} else {
format!("TABLES IN DATABASE `{}` ({}):", target_name, table_names.len())
};
let body: Vec<String> = table_names.iter().map(|t| format!(" {}", t)).collect();
Ok(format!("{}\n{}", header, body.join("\n")))
}
async fn list_tables_active_pg(state: &AppState, connection_id: &str) -> TuskResult<Vec<String>> {
let schemas = crate::commands::schema::list_schemas_core(state, connection_id).await?;
let mut all: Vec<String> = Vec::new();
for schema in &schemas {
let tables = list_tables_core(state, connection_id, schema).await?;
for t in tables {
all.push(format!("{}.{}", schema, t.name));
}
}
Ok(all)
}
async fn list_tables_other_pg(
app: &AppHandle,
state: &AppState,
connection_id: &str,
target_db: &str,
) -> TuskResult<Vec<String>> {
let cache_key = (connection_id.to_string(), target_db.to_string());
if let Some(hit) = state.tables_by_db_cache.read().await.get(&cache_key).cloned() {
if hit.cached_at.elapsed() < TOOL_CACHE_TTL {
return Ok(hit.value);
}
}
let config = load_connection_config(app, connection_id)?;
let url = config.connection_url_for_db(target_db);
let pool = PgPool::connect(&url).await.map_err(|e| {
TuskError::Custom(format!(
"Could not connect to database '{}' on this server: {}",
target_db, e
))
})?;
let rows = sqlx::query(
"SELECT table_schema, table_name FROM information_schema.tables \
WHERE table_schema NOT IN ('pg_catalog','information_schema','pg_toast','gp_toolkit') \
AND table_type = 'BASE TABLE' \
ORDER BY table_schema, table_name",
)
.fetch_all(&pool)
.await
.map_err(TuskError::Database)?;
pool.close().await;
let names: Vec<String> = rows
.iter()
.map(|r| format!("{}.{}", r.get::<String, _>(0), r.get::<String, _>(1)))
.collect();
state.tables_by_db_cache.write().await.insert(
cache_key,
CachedVec {
value: names.clone(),
cached_at: Instant::now(),
},
);
Ok(names)
}
async fn list_tables_clickhouse(
state: &AppState,
connection_id: &str,
target_db: &str,
) -> TuskResult<Vec<String>> {
let client = state.get_ch_client(connection_id).await?;
let escaped = target_db.replace('\\', "\\\\").replace('\'', "\\'");
let sql = format!(
"SELECT name FROM system.tables WHERE database = '{}' ORDER BY name",
escaped
);
let rows = client.fetch_objects(&sql).await?;
Ok(rows
.iter()
.filter_map(|r| r.get("name").and_then(|v| v.as_str()).map(String::from))
.collect())
}
// ---------------------------------------------------------------------------
// get_columns
// ---------------------------------------------------------------------------
pub async fn get_columns_tool(
state: &AppState,
connection_id: &str,
tables: &[String],
) -> TuskResult<String> {
if tables.is_empty() {
return Err(TuskError::Custom("get_columns requires at least one table.".into()));
}
if tables.len() > MAX_TABLES_PER_GET_COLUMNS {
return Err(TuskError::Custom(format!(
"Too many tables ({}); split into batches of ≤{}.",
tables.len(),
MAX_TABLES_PER_GET_COLUMNS
)));
}
let active_db = active_db_name(state, connection_id).await.unwrap_or_default();
// Normalise: accept "schema.table", "db.schema.table" (drop db if == active),
// and "table" (assume schema "public" for PG, or active DB for CH).
let parsed: Vec<(String, String, String)> = tables
.iter()
.map(|raw| normalise_table_ref(raw, &active_db))
.collect();
let flavor = state.get_flavor(connection_id).await;
if matches!(flavor, DbFlavor::ClickHouse) {
return get_columns_clickhouse(state, connection_id, &parsed).await;
}
get_columns_postgres(state, connection_id, &parsed).await
}
fn normalise_table_ref(raw: &str, active_db: &str) -> (String, String, String) {
// Returns (schema, table, original_input_for_diagnostics)
let trimmed = raw.trim().trim_matches('"').trim_matches('`');
let parts: Vec<&str> = trimmed.split('.').collect();
match parts.len() {
1 => ("public".to_string(), parts[0].to_string(), raw.to_string()),
2 => (parts[0].to_string(), parts[1].to_string(), raw.to_string()),
3 => {
// "db.schema.table" — drop db prefix when it matches active
let (db, schema, table) = (parts[0], parts[1], parts[2]);
if db == active_db {
(schema.to_string(), table.to_string(), raw.to_string())
} else {
// Different DB requested — let the caller surface a not-found warning.
// We still parse it as schema.table here.
(schema.to_string(), table.to_string(), raw.to_string())
}
}
_ => ("public".to_string(), trimmed.to_string(), raw.to_string()),
}
}
async fn get_columns_postgres(
state: &AppState,
connection_id: &str,
requested: &[(String, String, String)],
) -> TuskResult<String> {
let pool = state.get_pool(connection_id).await?;
let (col_res, fk_res, enum_res, tbl_comm_res, col_comm_res, unique_res) = tokio::join!(
fetch_columns(&pool),
fetch_foreign_keys_raw(&pool),
fetch_enum_types(&pool),
fetch_table_comments(&pool),
fetch_column_comments(&pool),
fetch_unique_constraints(&pool),
);
let all_cols = col_res?;
let fk_rows = fk_res?;
let enum_map = enum_res.unwrap_or_default();
let tbl_comments = tbl_comm_res.unwrap_or_default();
let col_comments = col_comm_res.unwrap_or_default();
let uniques = unique_res.unwrap_or_default();
// Build (schema, table) → Vec<ColumnInfo>
let mut by_table: BTreeMap<(String, String), Vec<ColumnInfo>> = BTreeMap::new();
for ci in &all_cols {
by_table
.entry((ci.schema.clone(), ci.table.clone()))
.or_default()
.push(ci.clone());
}
let mut fk_inline: HashMap<(String, String, String), String> = HashMap::new();
for fk in &fk_rows {
if fk.columns.len() == 1 && fk.ref_columns.len() == 1 {
fk_inline.insert(
(fk.schema.clone(), fk.table.clone(), fk.columns[0].clone()),
format!("{}.{}({})", fk.ref_schema, fk.ref_table, fk.ref_columns[0]),
);
}
}
let mut unique_map: HashMap<(String, String), Vec<String>> = HashMap::new();
for (schema, table, cols) in &uniques {
unique_map
.entry((schema.clone(), table.clone()))
.or_default()
.push(cols.join(", "));
}
let varchar_values: HashMap<(String, String, String), Vec<String>> = HashMap::new();
let jsonb_keys: HashMap<(String, String, String), Vec<String>> = HashMap::new();
let mut output: Vec<String> = Vec::new();
let mut not_found: Vec<String> = Vec::new();
for (schema, table, raw) in requested {
match by_table.get(&(schema.clone(), table.clone())) {
Some(cols) => {
let full_name = format!("{}.{}", schema, table);
format_table_block(
&full_name,
cols,
&tbl_comments,
&col_comments,
&fk_inline,
&enum_map,
&unique_map,
&varchar_values,
&jsonb_keys,
&mut output,
);
}
None => not_found.push(raw.clone()),
}
}
if !not_found.is_empty() {
let nearest = nearest_table_matches(&by_table, &not_found);
let header = format!(
"WARNING: tables not found: {}.{}",
not_found.join(", "),
if nearest.is_empty() {
String::new()
} else {
format!(" Nearest matches: {}.", nearest.join(", "))
}
);
output.insert(0, header);
output.insert(1, String::new());
}
let mut text = output.join("\n");
if text.len() > COLUMNS_TOOL_OUTPUT_CAP {
text.truncate(COLUMNS_TOOL_OUTPUT_CAP);
text.push_str("\n... (output truncated)");
}
Ok(text)
}
async fn get_columns_clickhouse(
state: &AppState,
connection_id: &str,
requested: &[(String, String, String)],
) -> TuskResult<String> {
let client = state.get_ch_client(connection_id).await?;
let active_db = client.database.clone();
let where_terms: Vec<String> = requested
.iter()
.map(|(schema, table, _)| {
// For CH, treat the parsed "schema" as the database name; if it equals
// a PG-conventional default ("public"), substitute with active CH database.
let dbn = if schema == "public" { active_db.clone() } else { schema.clone() };
format!(
"(database = '{}' AND name = '{}')",
dbn.replace('\'', "\\'"),
table.replace('\'', "\\'")
)
})
.collect();
let where_clause = where_terms.join(" OR ");
let sql = format!(
"SELECT database, table, name, type, default_expression, is_in_primary_key, comment, position \
FROM system.columns WHERE {} ORDER BY database, table, position",
where_clause
);
let rows = client.fetch_objects(&sql).await?;
// Group by (database, table)
let mut grouped: BTreeMap<(String, String), Vec<&serde_json::Map<String, serde_json::Value>>> =
BTreeMap::new();
for row in &rows {
let dbn = row.get("database").and_then(|v| v.as_str()).unwrap_or("").to_string();
let tbl = row.get("table").and_then(|v| v.as_str()).unwrap_or("").to_string();
grouped.entry((dbn, tbl)).or_default().push(row);
}
// Track which requested tables were found
let mut output = String::new();
let mut not_found: Vec<String> = Vec::new();
for (schema, table, raw) in requested {
let dbn = if schema == "public" { active_db.clone() } else { schema.clone() };
match grouped.get(&(dbn.clone(), table.clone())) {
Some(cols) => {
output.push_str(&format!("\nTABLE {}.{}\n", dbn, table));
for col in cols {
let name = col.get("name").and_then(|v| v.as_str()).unwrap_or("");
let dtype = col.get("type").and_then(|v| v.as_str()).unwrap_or("");
let is_pk = matches!(
col.get("is_in_primary_key"),
Some(serde_json::Value::Number(n)) if n.as_i64() == Some(1)
) || matches!(
col.get("is_in_primary_key"),
Some(serde_json::Value::String(s)) if s == "1"
);
let default = col.get("default_expression").and_then(|v| v.as_str()).unwrap_or("");
let comment = col.get("comment").and_then(|v| v.as_str()).unwrap_or("");
let mut line = format!(" {} {}", name, dtype);
if is_pk {
line.push_str(" [PK]");
}
if !default.is_empty() {
line.push_str(&format!(" DEFAULT {}", default));
}
if !comment.is_empty() {
line.push_str(&format!(" -- {}", comment));
}
output.push_str(&line);
output.push('\n');
}
}
None => not_found.push(raw.clone()),
}
}
let mut header = String::new();
if !not_found.is_empty() {
header.push_str(&format!(
"WARNING: tables not found: {}\n\n",
not_found.join(", ")
));
}
let mut combined = format!("{}{}", header, output.trim_start());
if combined.len() > COLUMNS_TOOL_OUTPUT_CAP {
combined.truncate(COLUMNS_TOOL_OUTPUT_CAP);
combined.push_str("\n... (output truncated)");
}
Ok(combined)
}
fn nearest_table_matches(
by_table: &BTreeMap<(String, String), Vec<ColumnInfo>>,
missing: &[String],
) -> Vec<String> {
let all: Vec<String> = by_table
.keys()
.map(|(s, t)| format!("{}.{}", s, t))
.collect();
let mut hints: Vec<String> = Vec::new();
for m in missing {
let needle = m.to_lowercase();
let mut candidates: Vec<&String> = all
.iter()
.filter(|n| {
let lower = n.to_lowercase();
lower.contains(&needle) || needle.contains(lower.split('.').last().unwrap_or(""))
})
.take(3)
.collect();
candidates.dedup();
for c in candidates {
if !hints.contains(c) {
hints.push(c.clone());
}
}
}
hints
}
// ---------------------------------------------------------------------------
// switch_database
// ---------------------------------------------------------------------------
pub async fn switch_database_tool(
app: &AppHandle,
state: &AppState,
connection_id: &str,
target_db: &str,
) -> TuskResult<String> {
let config = load_connection_config(app, connection_id)?;
// Verify target exists in cluster
let dbs = list_databases_core(state, connection_id).await?;
if !dbs.iter().any(|d| d == target_db) {
return Err(TuskError::Custom(format!(
"Database '{}' does not exist on this server. Available: {}",
target_db,
dbs.join(", ")
)));
}
switch_database_core(state, &config, target_db).await?;
Ok(format!("Switched active database to '{}'.", target_db))
}
// ---------------------------------------------------------------------------
// helpers
// ---------------------------------------------------------------------------
async fn active_db_name(state: &AppState, connection_id: &str) -> Option<String> {
let flavor = state.get_flavor(connection_id).await;
if matches!(flavor, DbFlavor::ClickHouse) {
return state
.get_ch_client(connection_id)
.await
.ok()
.map(|c| c.database.clone());
}
let pool = state.get_pool(connection_id).await.ok()?;
sqlx::query_scalar::<_, String>("SELECT current_database()")
.fetch_one(&pool)
.await
.ok()
}
// ---------------------------------------------------------------------------
// save_query / find_queries (chat v3 — F2)
// ---------------------------------------------------------------------------
const FIND_QUERIES_LIMIT: usize = 10;
const FIND_QUERIES_SQL_PREVIEW_CHARS: usize = 500;
pub async fn save_query_tool(
app: &AppHandle,
connection_id: &str,
name: &str,
sql: &str,
) -> TuskResult<String> {
let trimmed_name = name.trim();
let trimmed_sql = sql.trim();
if trimmed_name.is_empty() {
return Err(TuskError::Custom("save_query: name must not be empty".into()));
}
if trimmed_sql.is_empty() {
return Err(TuskError::Custom("save_query: sql must not be empty".into()));
}
let entry = SavedQuery {
id: uuid::Uuid::new_v4().to_string(),
name: trimmed_name.to_string(),
sql: trimmed_sql.to_string(),
connection_id: Some(connection_id.to_string()),
created_at: chrono::Utc::now().to_rfc3339(),
};
save_query_core(app, entry).await?;
Ok(format!("Saved query \"{}\" — visible in sidebar → Saved.", trimmed_name))
}
pub async fn find_queries_tool(
app: &AppHandle,
connection_id: &str,
text: &str,
) -> TuskResult<String> {
let trimmed = text.trim();
if trimmed.is_empty() {
return Err(TuskError::Custom("find_queries: text must not be empty".into()));
}
let all = list_saved_queries_core(app, Some(trimmed)).await?;
let matches: Vec<SavedQuery> = all
.into_iter()
.filter(|q| q.connection_id.as_deref() == Some(connection_id))
.take(FIND_QUERIES_LIMIT)
.collect();
if matches.is_empty() {
return Ok(format!(
"No saved queries match \"{}\" for this connection.",
trimmed
));
}
let mut out = format!(
"Saved queries matching \"{}\" ({}):",
trimmed,
matches.len()
);
for q in &matches {
let sql_preview: String = if q.sql.chars().count() > FIND_QUERIES_SQL_PREVIEW_CHARS {
let truncated: String = q.sql.chars().take(FIND_QUERIES_SQL_PREVIEW_CHARS).collect();
format!("{}", truncated)
} else {
q.sql.clone()
};
out.push_str(&format!(
"\n\n[{}] {}\n{}",
q.created_at, q.name, sql_preview
));
}
Ok(out)
}