diff --git a/crates/dbx-cli/src/commands.rs b/crates/dbx-cli/src/commands.rs index 349ff0e2dd..52e2324d9a 100644 --- a/crates/dbx-cli/src/commands.rs +++ b/crates/dbx-cli/src/commands.rs @@ -364,7 +364,9 @@ async fn safe_query(args: &[String]) -> CliEnvelope { }; let database = option_value(args, "--db").or(config.database.as_deref()).unwrap_or(""); - match dbx_core::query::execute_sql_statement_with_row_limit(&state, &config.id, database, sql, None, None, limit).await { + match dbx_core::query::execute_sql_statement_with_row_limit(&state, &config.id, database, sql, None, None, limit) + .await + { Ok(result) => ok( CliSource::Headless, serde_json::json!({ @@ -386,7 +388,7 @@ async fn handoff(args: &[String]) -> CliEnvelope { } let Some(conn) = required_option(args, "--conn") else { - return invalid_args("--conn is required"); + return fail(CliSource::Headless, CliErrorCode::ConnectionNotFound, "--conn is required", true); }; let Some(title) = required_option(args, "--title") else { return invalid_args("--title is required"); @@ -407,11 +409,15 @@ async fn handoff(args: &[String]) -> CliEnvelope { return invalid_args("SQL must not be empty"); } - let risk = dbx_core::sql_safety::risk_for_connection(&sql, conn, None); + let config = match find_connection(conn).await { + Ok(config) => config, + Err(err) => return err, + }; + let risk = dbx_core::sql_safety::risk_for_connection(&sql, &config.name, config.color.as_deref()); let item = dbx_core::handoff::HandoffItem::queued( - conn.to_string(), - conn.to_string(), - None, + config.id, + config.name, + config.database, title.to_string(), option_value(args, "--description").map(str::to_string), sql, @@ -593,16 +599,19 @@ mod tests { async fn create_large_sqlite_fixture(path: &std::path::Path, row_count: usize) { std::fs::File::create(path).unwrap(); let pool = dbx_core::db::sqlite::connect_path(&path.display().to_string()).await.unwrap(); - dbx_core::db::sqlite::execute_query(&pool, "CREATE TABLE numbers (id INTEGER PRIMARY KEY, value TEXT NOT NULL)") - .await - .unwrap(); + dbx_core::db::sqlite::execute_query( + &pool, + "CREATE TABLE numbers (id INTEGER PRIMARY KEY, value TEXT NOT NULL)", + ) + .await + .unwrap(); for id in 1..=row_count { dbx_core::db::sqlite::execute_query( &pool, &format!("INSERT INTO numbers (id, value) VALUES ({id}, 'value-{id}')"), ) - .await - .unwrap(); + .await + .unwrap(); } pool.close().await; } @@ -1183,7 +1192,7 @@ mod tests { assert_failure_code( dispatch(vec!["handoff".into(), "--title".into(), "Review".into(), "--sql".into(), "select 1".into()]) .await, - CliErrorCode::InternalError, + CliErrorCode::ConnectionNotFound, ); assert_failure_code( dispatch(vec!["handoff".into(), "--conn".into(), "local".into(), "--sql".into(), "select 1".into()]).await, @@ -1213,6 +1222,7 @@ mod tests { let sql_file = dir.path().join("query.sql"); std::fs::write(&sql_file, "select 1").unwrap(); std::env::set_var("DBX_APP_DATA_DIR", dir.path()); + seed_connections(dir.path(), &[mysql_fixture("local-id", "local", Some("#22c55e"), "local-secret")]).await; let env = dispatch(vec![ "handoff".into(), @@ -1236,6 +1246,49 @@ mod tests { std::env::remove_var("DBX_APP_DATA_DIR"); } + #[tokio::test] + async fn handoff_queues_with_resolved_connection_metadata_and_risk() { + let _guard = ENV_LOCK.lock().unwrap(); + let dir = tempfile::tempdir().unwrap(); + std::env::set_var("DBX_APP_DATA_DIR", dir.path()); + seed_connections(dir.path(), &[mysql_fixture("prod-id", "prod-main", Some("#ef4444"), "prod-secret")]).await; + + let env = dispatch(vec![ + "handoff".into(), + "--conn".into(), + "prod-id".into(), + "--title".into(), + "Review write".into(), + "--description".into(), + "generated by agent".into(), + "--sql".into(), + "UPDATE users SET active = 0 WHERE id = 1".into(), + "--format".into(), + "json".into(), + ]) + .await; + + match env { + CliEnvelope::Success { source, data, .. } => { + assert_eq!(source, CliSource::Headless); + assert_eq!(data["status"], "queued"); + } + CliEnvelope::Failure { error, .. } => panic!("expected queued handoff, got {error:?}"), + } + + let storage = dbx_core::storage::Storage::open(&dir.path().join("dbx.db")).await.unwrap(); + let queued = storage.load_pending_handoffs().await.unwrap(); + + assert_eq!(queued.len(), 1); + assert_eq!(queued[0].connection_id, "prod-id"); + assert_eq!(queued[0].connection_name, "prod-main"); + assert_eq!(queued[0].database.as_deref(), Some("app")); + assert_eq!(queued[0].risk_level, dbx_core::sql_safety::RiskLevel::High); + assert!(queued[0].is_production); + + std::env::remove_var("DBX_APP_DATA_DIR"); + } + #[tokio::test] async fn result_current_rejects_non_positive_and_over_limit_values() { for limit in ["0", "-1", "1001", "abc"] { diff --git a/crates/dbx-cli/src/main.rs b/crates/dbx-cli/src/main.rs index f34308a758..9a5a19356d 100644 --- a/crates/dbx-cli/src/main.rs +++ b/crates/dbx-cli/src/main.rs @@ -4,10 +4,7 @@ mod runtime_client; #[tokio::main] async fn main() { if let Err(err) = commands::run(std::env::args().skip(1).collect()).await { - println!( - "{}", - serde_json::to_string_pretty(&err).unwrap_or_else(|_| "{\"ok\":false}".to_string()) - ); + println!("{}", serde_json::to_string_pretty(&err).unwrap_or_else(|_| "{\"ok\":false}".to_string())); std::process::exit(1); } } diff --git a/crates/dbx-cli/src/runtime_client.rs b/crates/dbx-cli/src/runtime_client.rs index 81dda04c8e..effa954a0a 100644 --- a/crates/dbx-cli/src/runtime_client.rs +++ b/crates/dbx-cli/src/runtime_client.rs @@ -17,8 +17,7 @@ pub fn app_data_dir() -> PathBuf { return PathBuf::from(path); } - let home = std::env::var(if cfg!(windows) { "APPDATA" } else { "HOME" }) - .unwrap_or_else(|_| ".".to_string()); + let home = std::env::var(if cfg!(windows) { "APPDATA" } else { "HOME" }).unwrap_or_else(|_| ".".to_string()); if cfg!(target_os = "macos") { PathBuf::from(home).join("Library/Application Support/com.dbx.app") @@ -43,44 +42,30 @@ pub async fn get_json(path: &str) -> Result { let runtime = load_runtime().ok_or_else(|| "runtime unavailable".to_string())?; let url = runtime_url(&runtime, path, &[])?; - let response = reqwest::Client::new() - .get(url) - .bearer_auth(runtime.token) - .send() - .await - .map_err(|err| err.to_string())?; + let response = + reqwest::Client::new().get(url).bearer_auth(runtime.token).send().await.map_err(|err| err.to_string())?; let status = response.status(); if !status.is_success() { return Err(format!("runtime request failed with status {status}")); } - response - .json() - .await - .map_err(|err| err.to_string()) + response.json().await.map_err(|err| err.to_string()) } pub async fn get_json_with_query(path: &str, query: &[(&str, String)]) -> Result { let runtime = load_runtime().ok_or_else(|| "runtime unavailable".to_string())?; let url = runtime_url(&runtime, path, query)?; - let response = reqwest::Client::new() - .get(url) - .bearer_auth(runtime.token) - .send() - .await - .map_err(|err| err.to_string())?; + let response = + reqwest::Client::new().get(url).bearer_auth(runtime.token).send().await.map_err(|err| err.to_string())?; let status = response.status(); if !status.is_success() { return Err(format!("runtime request failed with status {status}")); } - response - .json() - .await - .map_err(|err| err.to_string()) + response.json().await.map_err(|err| err.to_string()) } pub async fn post_json(path: &str, body: serde_json::Value) -> Result { @@ -100,10 +85,7 @@ pub async fn post_json(path: &str, body: serde_json::Value) -> Result Result { diff --git a/crates/dbx-core/src/cli.rs b/crates/dbx-core/src/cli.rs index 534c4e1f41..6296e925a7 100644 --- a/crates/dbx-core/src/cli.rs +++ b/crates/dbx-core/src/cli.rs @@ -116,8 +116,12 @@ where let source = source.ok_or_else(|| de::Error::missing_field("source"))?; match (ok, has_data, has_error) { - (true, true, false) => Ok(CliEnvelope::Success { ok, source, data: data.expect("data presence was checked") }), - (false, false, true) => Ok(CliEnvelope::Failure { ok, source, error: error.expect("error presence was checked") }), + (true, true, false) => { + Ok(CliEnvelope::Success { ok, source, data: data.expect("data presence was checked") }) + } + (false, false, true) => { + Ok(CliEnvelope::Failure { ok, source, error: error.expect("error presence was checked") }) + } (true, _, _) => Err(de::Error::custom("ok=true requires data without error")), (false, _, _) => Err(de::Error::custom("ok=false requires error without data")), } diff --git a/crates/dbx-core/src/db/sqlite.rs b/crates/dbx-core/src/db/sqlite.rs index a826513238..f2d668820c 100644 --- a/crates/dbx-core/src/db/sqlite.rs +++ b/crates/dbx-core/src/db/sqlite.rs @@ -165,7 +165,11 @@ pub async fn execute_query(pool: &SqlitePool, sql: &str) -> Result Result { +pub async fn execute_query_with_row_limit( + pool: &SqlitePool, + sql: &str, + row_limit: usize, +) -> Result { let start = Instant::now(); let row_limit = row_limit.max(1); diff --git a/crates/dbx-core/src/storage.rs b/crates/dbx-core/src/storage.rs index 347caec01e..3950072916 100644 --- a/crates/dbx-core/src/storage.rs +++ b/crates/dbx-core/src/storage.rs @@ -4,7 +4,7 @@ use std::str::FromStr; use sqlx::sqlite::{SqliteConnectOptions, SqlitePool, SqlitePoolOptions}; use crate::ai::{AiChatMessage, AiConfig, AiConversation}; -use crate::handoff::HandoffItem; +use crate::handoff::{HandoffItem, HandoffStatus}; use crate::history::HistoryEntry; use crate::models::connection::ConnectionConfig; use crate::saved_sql::{SavedSqlFile, SavedSqlFolder, SavedSqlLibrary}; @@ -144,10 +144,11 @@ async fn ensure_history_columns(pool: &SqlitePool) -> Result<(), String> { } async fn ensure_handoffs_sequence(pool: &SqlitePool) -> Result<(), String> { - let (seq_columns,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM pragma_table_info('handoffs') WHERE name = 'seq'") - .fetch_one(pool) - .await - .map_err(|e| e.to_string())?; + let (seq_columns,): (i64,) = + sqlx::query_as("SELECT COUNT(*) FROM pragma_table_info('handoffs') WHERE name = 'seq'") + .fetch_one(pool) + .await + .map_err(|e| e.to_string())?; if seq_columns > 0 { return Ok(()); @@ -299,10 +300,7 @@ impl Storage { impl Storage { pub async fn save_handoff(&self, item: &HandoffItem) -> Result<(), String> { let json = serde_json::to_string(item).map_err(|e| e.to_string())?; - let status = serde_json::to_value(&item.status) - .ok() - .and_then(|value| value.as_str().map(str::to_string)) - .ok_or_else(|| "Failed to serialize handoff status".to_string())?; + let status = handoff_status_value(&item.status)?; sqlx::query( "INSERT INTO handoffs (id, payload_json, status, created_at) \ @@ -323,6 +321,32 @@ impl Storage { Ok(()) } + pub async fn update_handoff_status(&self, id: &str, status: HandoffStatus) -> Result { + let row: Option<(String,)> = sqlx::query_as("SELECT payload_json FROM handoffs WHERE id = ?") + .bind(id) + .fetch_optional(&self.db) + .await + .map_err(|e| e.to_string())?; + let Some((payload_json,)) = row else { + return Ok(false); + }; + + let mut item: HandoffItem = serde_json::from_str(&payload_json).map_err(|e| e.to_string())?; + item.status = status; + let updated_json = serde_json::to_string(&item).map_err(|e| e.to_string())?; + let status_value = handoff_status_value(&item.status)?; + + sqlx::query("UPDATE handoffs SET payload_json = ?, status = ? WHERE id = ?") + .bind(updated_json) + .bind(status_value) + .bind(id) + .execute(&self.db) + .await + .map_err(|e| e.to_string())?; + + Ok(true) + } + pub async fn load_pending_handoffs(&self) -> Result, String> { let rows: Vec<(String,)> = sqlx::query_as( "SELECT payload_json FROM handoffs \ @@ -337,6 +361,13 @@ impl Storage { } } +fn handoff_status_value(status: &HandoffStatus) -> Result { + serde_json::to_value(status) + .ok() + .and_then(|value| value.as_str().map(str::to_string)) + .ok_or_else(|| "Failed to serialize handoff status".to_string()) +} + // --------------------------------------------------------------------------- // AI Config // --------------------------------------------------------------------------- @@ -1057,10 +1088,8 @@ mod handoff_tests { async fn handoffs_table_has_stable_autoincrement_sequence() { let storage = open_temp_storage().await; - let columns: Vec<(String,)> = sqlx::query_as("SELECT name FROM pragma_table_info('handoffs')") - .fetch_all(&storage.db) - .await - .unwrap(); + let columns: Vec<(String,)> = + sqlx::query_as("SELECT name FROM pragma_table_info('handoffs')").fetch_all(&storage.db).await.unwrap(); assert!(columns.iter().any(|(name,)| name == "seq")); } @@ -1091,6 +1120,25 @@ mod handoff_tests { assert_eq!(loaded.len(), 1); assert_eq!(loaded[0].id, queued.id); } + + #[tokio::test] + async fn update_handoff_status_updates_payload_and_pending_visibility() { + let storage = open_temp_storage().await; + let item = queued_handoff("review me"); + let missing = uuid::Uuid::new_v4().to_string(); + + storage.save_handoff(&item).await.unwrap(); + + assert!(storage.update_handoff_status(&item.id, HandoffStatus::Shown).await.unwrap()); + assert!(!storage.update_handoff_status(&missing, HandoffStatus::Rejected).await.unwrap()); + + let shown = storage.load_pending_handoffs().await.unwrap(); + assert_eq!(shown.len(), 1); + assert_eq!(shown[0].status, HandoffStatus::Shown); + + assert!(storage.update_handoff_status(&item.id, HandoffStatus::Rejected).await.unwrap()); + assert!(storage.load_pending_handoffs().await.unwrap().is_empty()); + } } // --------------------------------------------------------------------------- diff --git a/src-tauri/src/commands/agent_runtime.rs b/src-tauri/src/commands/agent_runtime.rs index 8c1d093b2b..d88a2ef287 100644 --- a/src-tauri/src/commands/agent_runtime.rs +++ b/src-tauri/src/commands/agent_runtime.rs @@ -9,6 +9,7 @@ use tokio::net::{TcpListener, TcpStream}; use tokio::sync::{oneshot, RwLock}; use super::connection::AppState; +use dbx_core::handoff::{HandoffItem, HandoffStatus}; const BIND_ADDR: &str = "127.0.0.1:0"; const DISCOVERY_FILE: &str = "agent-runtime.json"; @@ -97,10 +98,28 @@ pub async fn agent_runtime_load_handoffs( runtime: tauri::State<'_, AgentRuntimeServer>, ) -> Result, String> { let mut items = app_state.storage.load_pending_handoffs().await?; - items.extend(runtime.state().handoffs.read().await.iter().cloned()); + items.extend(pending_runtime_handoffs(runtime.state()).await); Ok(items) } +#[tauri::command] +pub async fn agent_runtime_mark_handoff_shown( + app_state: tauri::State<'_, Arc>, + runtime: tauri::State<'_, AgentRuntimeServer>, + id: String, +) -> Result { + update_handoff_status(app_state.inner().as_ref(), runtime.state(), &id, HandoffStatus::Shown).await +} + +#[tauri::command] +pub async fn agent_runtime_reject_handoff( + app_state: tauri::State<'_, Arc>, + runtime: tauri::State<'_, AgentRuntimeServer>, + id: String, +) -> Result { + update_handoff_status(app_state.inner().as_ref(), runtime.state(), &id, HandoffStatus::Rejected).await +} + pub fn start(app: AppHandle) -> AgentRuntimeServer { let token = uuid::Uuid::new_v4().to_string(); let state = AgentRuntimeState { @@ -306,10 +325,7 @@ async fn route_request(first_line: &str, body: &str, state: &AgentRuntimeState) if let Some(limit) = query_limit(first_line) { truncate_result_rows(&mut body, limit); } - return RuntimeResponse { - status: "200 OK", - body, - }; + return RuntimeResponse { status: "200 OK", body }; } if first_line.starts_with("POST /handoff ") { @@ -331,6 +347,37 @@ async fn route_request(first_line: &str, body: &str, state: &AgentRuntimeState) RuntimeResponse { status: "404 Not Found", body: serde_json::json!({"error": "not found"}) } } +async fn update_handoff_status( + app_state: &AppState, + runtime: &AgentRuntimeState, + id: &str, + status: HandoffStatus, +) -> Result { + let stored = app_state.storage.update_handoff_status(id, status.clone()).await?; + let runtime_updated = update_runtime_handoff_status(runtime, id, status).await; + Ok(stored || runtime_updated) +} + +async fn update_runtime_handoff_status(state: &AgentRuntimeState, id: &str, status: HandoffStatus) -> bool { + let mut handoffs = state.handoffs.write().await; + if let Some(item) = handoffs.iter_mut().find(|item| item.id == id) { + item.status = status; + return true; + } + false +} + +async fn pending_runtime_handoffs(state: &AgentRuntimeState) -> Vec { + state + .handoffs + .read() + .await + .iter() + .filter(|item| matches!(item.status, HandoffStatus::Queued | HandoffStatus::Shown)) + .cloned() + .collect() +} + fn query_limit(first_line: &str) -> Option { let target = first_line.split_whitespace().nth(1)?; let query = target.split_once('?')?.1; @@ -563,6 +610,30 @@ mod tests { assert!(response.starts_with("HTTP/1.1 413 Payload Too Large"), "{response}"); } + #[tokio::test] + async fn runtime_handoff_status_updates_filter_pending_items() { + let state = runtime_state(); + let item = dbx_core::handoff::HandoffItem::queued( + "conn-1".to_string(), + "Local".to_string(), + Some("main".to_string()), + "Review SQL".to_string(), + None, + "update users set name = 'a'".to_string(), + dbx_core::sql_safety::OperationClass::Write, + dbx_core::sql_safety::RiskLevel::Medium, + false, + ); + let id = item.id.clone(); + state.handoffs.write().await.push(item); + + assert!(update_runtime_handoff_status(&state, &id, dbx_core::handoff::HandoffStatus::Shown).await); + assert_eq!(pending_runtime_handoffs(&state).await[0].status, dbx_core::handoff::HandoffStatus::Shown); + + assert!(update_runtime_handoff_status(&state, &id, dbx_core::handoff::HandoffStatus::Rejected).await); + assert!(pending_runtime_handoffs(&state).await.is_empty()); + } + #[test] fn discovery_file_is_owner_only_and_removed_on_cleanup() { let dir = std::env::temp_dir().join(format!("dbx-agent-runtime-test-{}", uuid::Uuid::new_v4())); diff --git a/src-tauri/src/commands/mod.rs b/src-tauri/src/commands/mod.rs index 2040e4ca0b..69f61492d8 100644 --- a/src-tauri/src/commands/mod.rs +++ b/src-tauri/src/commands/mod.rs @@ -1,5 +1,5 @@ -pub mod ai; pub mod agent_runtime; +pub mod ai; pub mod connection; #[allow(dead_code, unused_imports)] mod connection_secrets; diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 7ce3f2cda2..cca1d6a8e5 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -61,6 +61,8 @@ pub fn run() { .invoke_handler(tauri::generate_handler![ commands::agent_runtime::agent_runtime_update_snapshot, commands::agent_runtime::agent_runtime_load_handoffs, + commands::agent_runtime::agent_runtime_mark_handoff_shown, + commands::agent_runtime::agent_runtime_reject_handoff, commands::ai::ai_complete, commands::ai::ai_stream, commands::ai::ai_cancel_stream, diff --git a/src/components/agent/AgentHandoffDialog.vue b/src/components/agent/AgentHandoffDialog.vue new file mode 100644 index 0000000000..610cd143cb --- /dev/null +++ b/src/components/agent/AgentHandoffDialog.vue @@ -0,0 +1,78 @@ + + + diff --git a/src/components/layout/AppDialogs.vue b/src/components/layout/AppDialogs.vue index 0061e9d12f..159fc25f08 100644 --- a/src/components/layout/AppDialogs.vue +++ b/src/components/layout/AppDialogs.vue @@ -4,6 +4,7 @@ import { useI18n } from "vue-i18n"; import { Button } from "@/components/ui/button"; import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; import ConnectionDialog from "@/components/connection/ConnectionDialog.vue"; +import AgentHandoffDialog from "@/components/agent/AgentHandoffDialog.vue"; import EditorSettingsDialog from "@/components/editor/EditorSettingsDialog.vue"; import DangerConfirmDialog from "@/components/editor/DangerConfirmDialog.vue"; const DataTransferDialog = defineAsyncComponent(() => import("@/components/transfer/DataTransferDialog.vue")); @@ -92,6 +93,7 @@ watch(