From b0d4575455366daaf9e1734c163c9cee2f33eba6 Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 15:34:16 +0800 Subject: [PATCH 01/40] chore: ignore worktree directory --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 0cd799edb..ef8705469 100644 --- a/.gitignore +++ b/.gitignore @@ -40,6 +40,7 @@ Thumbs.db # Temp tmp/ +.worktrees/ # Logs *.log From 4991acf79af00869109fb78242a0568bcec8c442 Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 15:45:24 +0800 Subject: [PATCH 02/40] docs: add dbx cli runtime implementation plan --- .../plans/2026-05-10-dbx-cli-runtime.md | 1537 +++++++++++++++++ 1 file changed, 1537 insertions(+) create mode 100644 docs/superpowers/plans/2026-05-10-dbx-cli-runtime.md diff --git a/docs/superpowers/plans/2026-05-10-dbx-cli-runtime.md b/docs/superpowers/plans/2026-05-10-dbx-cli-runtime.md new file mode 100644 index 000000000..961f99829 --- /dev/null +++ b/docs/superpowers/plans/2026-05-10-dbx-cli-runtime.md @@ -0,0 +1,1537 @@ +# DBX CLI Runtime Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add a full DBX CLI with independent GUI Runtime support for the 8 specified commands while preserving the existing MCP package and behavior. + +**Architecture:** Add a new `dbx-cli` Rust binary crate that first attempts GUI Runtime discovery, then falls back to headless mode. Shared CLI data models, JSON envelope, SQL safety classification, schema snapshot, risk metadata, and handoff queue live in `dbx-core`. The desktop app starts a separate authenticated localhost runtime and the Vue UI synchronizes active context, selection, result samples, and handoff review state into that runtime. + +**Tech Stack:** Rust workspace, `dbx-core`, Tauri v2, Vue 3/Pinia, SQLite storage via `sqlx`, existing DBX schema/query/SSH core, Node test runner for frontend utilities, Cargo tests for Rust. + +--- + +## Non-Negotiable Constraints + +- Preserve all existing `mcp/` files, package behavior, tool names, and bridge endpoints. +- Do not remove or rename current Tauri commands. +- Do not add CLI commands that create, edit, or delete DBX connections. +- All CLI agent-facing commands return the envelope shape `{ ok, source, data }` or `{ ok, source, error }`. +- `safe-query` directly executes only READ SQL; WRITE and DDL return structured blocking errors and point to `handoff`. +- `selection` and `result current` return real GUI state when runtime is available, not a placeholder. +- `handoff` sends to GUI runtime when running and writes queued pending records when headless. + +## File Structure + +- Create `crates/dbx-cli/Cargo.toml`: binary crate configuration. +- Create `crates/dbx-cli/src/main.rs`: argument parsing, command dispatch, runtime/headless selection. +- Create `crates/dbx-cli/src/commands.rs`: CLI command handlers. +- Create `crates/dbx-cli/src/runtime_client.rs`: discovery file loading, token-authenticated HTTP calls. +- Create `crates/dbx-core/src/cli.rs`: envelope, error codes, CLI DTOs, app data path helper. +- Create `crates/dbx-core/src/sql_safety.rs`: SQL classification and risk metadata. +- Create `crates/dbx-core/src/schema_snapshot.rs`: normalized schema snapshot orchestration. +- Create `crates/dbx-core/src/handoff.rs`: handoff model and queued storage helpers. +- Modify `crates/dbx-core/src/lib.rs`: export new modules. +- Modify `crates/dbx-core/src/storage.rs`: add pending handoff table migration and load/save helpers. +- Modify root `Cargo.toml`: add `crates/dbx-cli` workspace member. +- Modify `src-tauri/src/lib.rs`: start independent agent runtime and register runtime commands. +- Create `src-tauri/src/commands/agent_runtime.rs`: authenticated runtime server and shared runtime state. +- Modify `src-tauri/src/commands/mod.rs`: expose `agent_runtime`. +- Modify `src/lib/api.ts` and `src/lib/tauri.ts`: add runtime state sync and handoff API wrappers. +- Create `src/stores/agentRuntimeStore.ts`: collect current UI context and push snapshots. +- Create `src/components/agent/AgentHandoffDialog.vue`: review queued/runtime handoffs. +- Modify `src/components/layout/AppDialogs.vue`: mount handoff dialog. +- Modify `src/stores/queryStore.ts`: notify runtime store on active tab/result changes. +- Modify `src/components/grid/DataGrid.vue`: notify runtime store on grid selection changes. +- Modify `src/composables/useTauriEvents.ts`: keep existing MCP listeners unchanged. +- Add Rust tests under `crates/dbx-core/src/sql_safety.rs` and `crates/dbx-core/src/cli.rs`. +- Add frontend tests under `tests/agentRuntimeStore.test.ts` for payload shaping. + +--- + +### Task 1: Workspace and CLI Crate Skeleton + +**Files:** +- Modify: `/Users/bytedance/open_source_poj/dbx/Cargo.toml` +- Create: `/Users/bytedance/open_source_poj/dbx/crates/dbx-cli/Cargo.toml` +- Create: `/Users/bytedance/open_source_poj/dbx/crates/dbx-cli/src/main.rs` + +- [ ] **Step 1: Add workspace member** + +Change the workspace members in `/Users/bytedance/open_source_poj/dbx/Cargo.toml` to include the CLI crate: + +```toml +[workspace] +resolver = "2" +members = ["src-tauri", "crates/dbx-core", "crates/dbx-cli", "src-web"] +``` + +- [ ] **Step 2: Create CLI crate manifest** + +Create `/Users/bytedance/open_source_poj/dbx/crates/dbx-cli/Cargo.toml`: + +```toml +[package] +name = "dbx-cli" +version = "0.5.2" +edition = "2021" + +[[bin]] +name = "dbx" +path = "src/main.rs" + +[dependencies] +dbx-core = { path = "../dbx-core" } +chrono = { version = "0.4", features = ["serde"] } +serde = { version = "1.0", features = ["derive"] } +serde_json = "1.0" +tokio = { version = "1", features = ["full"] } +reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] } +uuid = { version = "1", features = ["v4", "serde"] } +``` + +- [ ] **Step 3: Create placeholder CLI entry** + +Create `/Users/bytedance/open_source_poj/dbx/crates/dbx-cli/src/main.rs`: + +```rust +mod commands; +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())); + std::process::exit(1); + } +} +``` + +- [ ] **Step 4: Verify workspace recognizes the binary** + +Run: + +```bash +cargo metadata --format-version 1 --no-deps +``` + +Expected: output contains `"name":"dbx-cli"` and `"name":"dbx"`. + +--- + +### Task 2: Core CLI Envelope and Error DTOs + +**Files:** +- Create: `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/cli.rs` +- Modify: `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/lib.rs` + +- [ ] **Step 1: Add envelope and error types** + +Create `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/cli.rs`: + +```rust +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "kebab-case")] +pub enum CliSource { + GuiRuntime, + Headless, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "SCREAMING_SNAKE_CASE")] +pub enum CliErrorCode { + GuiRuntimeRequired, + ConnectionNotFound, + AmbiguousConnection, + SecretUnavailable, + SshTunnelFailed, + QueryClassificationFailed, + HandoffRequired, + DdlBlocked, + ProductionWriteBlocked, + UnsupportedDatabaseType, + Timeout, + InternalError, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CliError { + pub code: CliErrorCode, + pub message: String, + pub recoverable: bool, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(untagged)] +pub enum CliEnvelope { + Success { ok: bool, source: CliSource, data: T }, + Failure { ok: bool, source: CliSource, error: CliError }, +} + +pub fn ok(source: CliSource, data: T) -> CliEnvelope { + CliEnvelope::Success { ok: true, source, data } +} + +pub fn fail(source: CliSource, code: CliErrorCode, message: impl Into, recoverable: bool) -> CliEnvelope { + CliEnvelope::Failure { ok: false, source, error: CliError { code, message: message.into(), recoverable } } +} +``` + +- [ ] **Step 2: Export module** + +Add this line to `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/lib.rs`: + +```rust +pub mod cli; +``` + +- [ ] **Step 3: Add envelope unit test** + +Append tests in `cli.rs`: + +```rust +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn serializes_success_source_as_kebab_case() { + let env = ok(CliSource::GuiRuntime, serde_json::json!({"value": 1})); + let json = serde_json::to_string(&env).unwrap(); + assert!(json.contains("\"ok\":true")); + assert!(json.contains("\"source\":\"gui-runtime\"")); + } + + #[test] + fn serializes_error_code_as_screaming_snake_case() { + let env: CliEnvelope<()> = fail(CliSource::Headless, CliErrorCode::GuiRuntimeRequired, "runtime needed", true); + let json = serde_json::to_string(&env).unwrap(); + assert!(json.contains("\"GUI_RUNTIME_REQUIRED\"")); + } +} +``` + +- [ ] **Step 4: Run targeted test** + +Run: + +```bash +cargo test -p dbx-core cli::tests +``` + +Expected: both tests pass. + +--- + +### Task 3: SQL Safety and Risk Classification + +**Files:** +- Create: `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/sql_safety.rs` +- Modify: `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/lib.rs` + +- [ ] **Step 1: Implement classifier** + +Create `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/sql_safety.rs`: + +```rust +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "lowercase")] +pub enum OperationClass { + Read, + Write, + Ddl, + Unknown, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "lowercase")] +pub enum RiskLevel { + Low, + Medium, + High, + Critical, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct RiskMetadata { + pub operation_class: OperationClass, + pub risk_level: RiskLevel, + pub is_production: bool, + pub production_reason: Option, + pub first_token: Option, +} + +pub fn classify_sql(sql: &str) -> OperationClass { + let token = first_executable_token(sql).map(|s| s.to_ascii_uppercase()); + match token.as_deref() { + Some("SELECT" | "SHOW" | "DESCRIBE" | "EXPLAIN" | "WITH") => OperationClass::Read, + Some("INSERT" | "UPDATE" | "DELETE" | "MERGE" | "REPLACE") => OperationClass::Write, + Some("CREATE" | "ALTER" | "DROP" | "TRUNCATE" | "RENAME") => OperationClass::Ddl, + _ => OperationClass::Unknown, + } +} + +pub fn risk_for(sql: &str, connection_name: &str, color: Option<&str>) -> RiskMetadata { + let operation_class = classify_sql(sql); + let (is_production, production_reason) = production_signal(connection_name, color); + let risk_level = match (operation_class, is_production) { + (OperationClass::Read, _) => RiskLevel::Low, + (OperationClass::Write, false) => RiskLevel::Medium, + (OperationClass::Write, true) => RiskLevel::High, + (OperationClass::Ddl, _) => RiskLevel::Critical, + (OperationClass::Unknown, _) => RiskLevel::High, + }; + RiskMetadata { + operation_class, + risk_level, + is_production, + production_reason, + first_token: first_executable_token(sql).map(str::to_string), + } +} + +fn production_signal(connection_name: &str, color: Option<&str>) -> (bool, Option) { + if matches!(color, Some("#ef4444")) { + return (true, Some("red connection color".to_string())); + } + let name = connection_name.to_ascii_lowercase(); + if ["prod", "production", "live"].iter().any(|needle| name.contains(needle)) { + return (true, Some("connection name fallback".to_string())); + } + (false, None) +} + +fn first_executable_token(sql: &str) -> Option<&str> { + let bytes = sql.as_bytes(); + let mut i = 0; + while i < bytes.len() { + while i < bytes.len() && bytes[i].is_ascii_whitespace() { + i += 1; + } + if i + 1 < bytes.len() && bytes[i] == b'-' && bytes[i + 1] == b'-' { + i += 2; + while i < bytes.len() && bytes[i] != b'\n' { + i += 1; + } + continue; + } + if i + 1 < bytes.len() && bytes[i] == b'/' && bytes[i + 1] == b'*' { + i += 2; + while i + 1 < bytes.len() && !(bytes[i] == b'*' && bytes[i + 1] == b'/') { + i += 1; + } + i = (i + 2).min(bytes.len()); + continue; + } + break; + } + let start = i; + while i < bytes.len() && (bytes[i].is_ascii_alphabetic() || bytes[i] == b'_') { + i += 1; + } + (i > start).then_some(&sql[start..i]) +} +``` + +- [ ] **Step 2: Export module** + +Add this line to `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/lib.rs`: + +```rust +pub mod sql_safety; +``` + +- [ ] **Step 3: Add classifier tests** + +Append tests to `sql_safety.rs`: + +```rust +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn comments_do_not_hide_read_token() { + assert_eq!(classify_sql("-- comment\nSELECT 1"), OperationClass::Read); + assert_eq!(classify_sql("/* DROP TABLE x */ SELECT 1"), OperationClass::Read); + } + + #[test] + fn classifies_write_and_ddl() { + assert_eq!(classify_sql("update users set name = 'a'"), OperationClass::Write); + assert_eq!(classify_sql("DROP TABLE users"), OperationClass::Ddl); + } + + #[test] + fn red_color_marks_production() { + let risk = risk_for("orders", "prod-main", Some("#ef4444")); + assert!(risk.is_production); + assert_eq!(risk.risk_level, RiskLevel::Low); + } +} +``` + +- [ ] **Step 4: Run targeted tests** + +Run: + +```bash +cargo test -p dbx-core sql_safety::tests +``` + +Expected: all classifier tests pass. + +--- + +### Task 4: Handoff Storage Model + +**Files:** +- Create: `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/handoff.rs` +- Modify: `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/lib.rs` +- Modify: `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/storage.rs` + +- [ ] **Step 1: Add handoff types** + +Create `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/handoff.rs`: + +```rust +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +use crate::sql_safety::{OperationClass, RiskLevel}; + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "lowercase")] +pub enum HandoffStatus { + Queued, + Shown, + Approved, + Rejected, + Executed, + Failed, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct HandoffItem { + pub id: String, + pub created_at: DateTime, + pub created_by: String, + pub connection_name: String, + pub database: Option, + pub title: String, + pub description: Option, + pub sql: String, + pub operation_class: OperationClass, + pub risk_level: RiskLevel, + pub is_production: bool, + pub status: HandoffStatus, + pub result_summary: Option, + pub error: Option, +} + +impl HandoffItem { + pub fn queued( + connection_name: String, + database: Option, + title: String, + description: Option, + sql: String, + operation_class: OperationClass, + risk_level: RiskLevel, + is_production: bool, + ) -> Self { + Self { + id: Uuid::new_v4().to_string(), + created_at: Utc::now(), + created_by: "dbx-cli".to_string(), + connection_name, + database, + title, + description, + sql, + operation_class, + risk_level, + is_production, + status: HandoffStatus::Queued, + result_summary: None, + error: None, + } + } +} +``` + +- [ ] **Step 2: Export module** + +Add this line to `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/lib.rs`: + +```rust +pub mod handoff; +``` + +- [ ] **Step 3: Add storage table** + +Add this SQL statement to `SCHEMA_STATEMENTS` in `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/storage.rs`: + +```rust +"CREATE TABLE IF NOT EXISTS handoffs ( + id TEXT PRIMARY KEY, + payload_json TEXT NOT NULL, + status TEXT NOT NULL, + created_at TEXT NOT NULL +)", +``` + +- [ ] **Step 4: Add storage methods** + +Append this impl block near the storage sections in `storage.rs`: + +```rust +impl Storage { + pub async fn save_handoff(&self, item: &crate::handoff::HandoffItem) -> Result<(), String> { + let json = serde_json::to_string(item).map_err(|e| e.to_string())?; + sqlx::query( + "INSERT OR REPLACE INTO handoffs (id, payload_json, status, created_at) VALUES (?, ?, ?, ?)", + ) + .bind(&item.id) + .bind(json) + .bind(format!("{:?}", item.status).to_lowercase()) + .bind(item.created_at.to_rfc3339()) + .execute(&self.db) + .await + .map_err(|e| e.to_string())?; + Ok(()) + } + + pub async fn load_pending_handoffs(&self) -> Result, String> { + let rows: Vec<(String,)> = sqlx::query_as( + "SELECT payload_json FROM handoffs WHERE status IN ('queued', 'shown') ORDER BY created_at ASC", + ) + .fetch_all(&self.db) + .await + .map_err(|e| e.to_string())?; + rows.into_iter() + .map(|(json,)| serde_json::from_str(&json).map_err(|e| e.to_string())) + .collect() + } +} +``` + +- [ ] **Step 5: Run core tests** + +Run: + +```bash +cargo test -p dbx-core +``` + +Expected: existing tests pass and storage compiles with new methods. + +--- + +### Task 5: Schema Snapshot Orchestration + +**Files:** +- Create: `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/schema_snapshot.rs` +- Modify: `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/lib.rs` + +- [ ] **Step 1: Add snapshot DTOs and function** + +Create `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/schema_snapshot.rs`: + +```rust +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; + +use crate::connection::AppState; +use crate::models::connection::DatabaseType; +use crate::{schema, types}; + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct TableSnapshot { + pub name: String, + pub table_type: String, + pub columns: Vec, + pub indexes: Vec, + pub foreign_keys: Vec, + pub triggers: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct SchemaSnapshot { + pub connection_id: String, + pub connection_name: String, + pub database: Option, + pub database_type: DatabaseType, + pub driver_profile: Option, + pub captured_at: DateTime, + pub databases: Vec, + pub schemas: Vec, + pub tables: Vec, +} + +pub async fn snapshot( + state: &AppState, + connection_id: &str, + database: Option<&str>, + schema_name: Option<&str>, +) -> Result { + let config = { + let configs = state.configs.lock().await; + configs.get(connection_id).cloned().ok_or("Connection config not found")? + }; + let db = database.or(config.database.as_deref()).unwrap_or_default(); + let schemas = if db.is_empty() { + Vec::new() + } else { + schema::list_schemas_core(state, connection_id, db).await.unwrap_or_default() + }; + let effective_schema = schema_name.or_else(|| schemas.first().map(String::as_str)).unwrap_or(""); + let databases = schema::list_databases_core(state, connection_id).await.unwrap_or_default(); + let table_infos = if db.is_empty() { + Vec::new() + } else { + schema::list_tables_core(state, connection_id, db, effective_schema).await.unwrap_or_default() + }; + let mut tables = Vec::new(); + for table in table_infos { + let columns = schema::get_columns_core(state, connection_id, db, effective_schema, &table.name).await.unwrap_or_default(); + let indexes = schema::list_indexes_core(state, connection_id, db, effective_schema, &table.name).await.unwrap_or_default(); + let foreign_keys = schema::list_foreign_keys_core(state, connection_id, db, effective_schema, &table.name).await.unwrap_or_default(); + let triggers = schema::list_triggers_core(state, connection_id, db, effective_schema, &table.name).await.unwrap_or_default(); + tables.push(TableSnapshot { name: table.name, table_type: table.table_type, columns, indexes, foreign_keys, triggers }); + } + Ok(SchemaSnapshot { + connection_id: config.id, + connection_name: config.name, + database: (!db.is_empty()).then(|| db.to_string()), + database_type: config.db_type, + driver_profile: config.driver_profile, + captured_at: Utc::now(), + databases, + schemas, + tables, + }) +} +``` + +- [ ] **Step 2: Export module** + +Add this line to `/Users/bytedance/open_source_poj/dbx/crates/dbx-core/src/lib.rs`: + +```rust +pub mod schema_snapshot; +``` + +- [ ] **Step 3: Compile core** + +Run: + +```bash +cargo check -p dbx-core +``` + +Expected: `dbx-core` compiles. + +--- + +### Task 6: Runtime Discovery and CLI Command Dispatch + +**Files:** +- Create: `/Users/bytedance/open_source_poj/dbx/crates/dbx-cli/src/runtime_client.rs` +- Create: `/Users/bytedance/open_source_poj/dbx/crates/dbx-cli/src/commands.rs` + +- [ ] **Step 1: Add runtime client** + +Create `/Users/bytedance/open_source_poj/dbx/crates/dbx-cli/src/runtime_client.rs`: + +```rust +use serde::{Deserialize, Serialize}; +use std::path::PathBuf; + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct RuntimeDiscovery { + pub port: u16, + pub token: String, +} + +pub fn app_data_dir() -> PathBuf { + 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") + } else if cfg!(windows) { + PathBuf::from(home).join("com.dbx.app") + } else { + PathBuf::from(home).join(".config/com.dbx.app") + } +} + +pub fn load_runtime() -> Option { + let path = app_data_dir().join("agent-runtime.json"); + let json = std::fs::read_to_string(path).ok()?; + serde_json::from_str(&json).ok() +} + +pub async fn get_json(path: &str) -> Result { + let runtime = load_runtime().ok_or("runtime unavailable")?; + let url = format!("http://127.0.0.1:{}{}", runtime.port, path); + reqwest::Client::new() + .get(url) + .bearer_auth(runtime.token) + .send() + .await + .map_err(|e| e.to_string())? + .json() + .await + .map_err(|e| e.to_string()) +} + +pub async fn post_json(path: &str, body: serde_json::Value) -> Result { + let runtime = load_runtime().ok_or("runtime unavailable")?; + let url = format!("http://127.0.0.1:{}{}", runtime.port, path); + reqwest::Client::new() + .post(url) + .bearer_auth(runtime.token) + .json(&body) + .send() + .await + .map_err(|e| e.to_string())? + .json() + .await + .map_err(|e| e.to_string()) +} +``` + +- [ ] **Step 2: Add command dispatcher skeleton** + +Create `/Users/bytedance/open_source_poj/dbx/crates/dbx-cli/src/commands.rs` with handlers for all 8 commands: + +```rust +use dbx_core::cli::{fail, ok, CliEnvelope, CliErrorCode, CliSource}; + +pub async fn run(args: Vec) -> Result<(), CliEnvelope<()>> { + let output = match args.as_slice() { + [cmd, rest @ ..] if cmd == "context" => context(rest).await, + [cmd, sub, rest @ ..] if cmd == "conn" && sub == "list" => conn_list(rest).await, + [cmd, sub, name, rest @ ..] if cmd == "conn" && sub == "show" => conn_show(name, rest).await, + [cmd, sub, rest @ ..] if cmd == "schema" && sub == "snapshot" => schema_snapshot(rest).await, + [cmd, rest @ ..] if cmd == "safe-query" => safe_query(rest).await, + [cmd, rest @ ..] if cmd == "handoff" => handoff(rest).await, + [cmd, rest @ ..] if cmd == "selection" => selection(rest).await, + [cmd, sub, rest @ ..] if cmd == "result" && sub == "current" => result_current(rest).await, + _ => fail(CliSource::Headless, CliErrorCode::InternalError, "Unknown command", false), + }; + println!("{}", serde_json::to_string_pretty(&output).unwrap()); + if matches!(output, CliEnvelope::Failure { .. }) { + std::process::exit(1); + } + Ok(()) +} + +async fn context(_args: &[String]) -> CliEnvelope { + match crate::runtime_client::get_json("/context").await { + Ok(data) => ok(CliSource::GuiRuntime, data), + Err(_) => ok(CliSource::Headless, serde_json::json!({ "runtime": "headless" })), + } +} + +async fn conn_list(_args: &[String]) -> CliEnvelope { + ok(CliSource::Headless, serde_json::json!({ "connections": [] })) +} + +async fn conn_show(_name: &str, _args: &[String]) -> CliEnvelope { + ok(CliSource::Headless, serde_json::json!({})) +} + +async fn schema_snapshot(_args: &[String]) -> CliEnvelope { + fail(CliSource::Headless, CliErrorCode::ConnectionNotFound, "Connection is required", true) +} + +async fn safe_query(_args: &[String]) -> CliEnvelope { + fail(CliSource::Headless, CliErrorCode::QueryClassificationFailed, "SQL is required", true) +} + +async fn handoff(_args: &[String]) -> CliEnvelope { + fail(CliSource::Headless, CliErrorCode::InternalError, "Handoff payload is required", true) +} + +async fn selection(_args: &[String]) -> CliEnvelope { + match crate::runtime_client::get_json("/selection").await { + Ok(data) => ok(CliSource::GuiRuntime, data), + Err(_) => fail(CliSource::Headless, CliErrorCode::GuiRuntimeRequired, "dbx selection requires DBX GUI runtime.", true), + } +} + +async fn result_current(args: &[String]) -> CliEnvelope { + let limit = option_value(args, "--limit").unwrap_or("50"); + match crate::runtime_client::get_json(&format!("/result/current?limit={limit}")).await { + Ok(data) => ok(CliSource::GuiRuntime, data), + Err(_) => fail(CliSource::Headless, CliErrorCode::GuiRuntimeRequired, "dbx result current requires DBX GUI runtime.", true), + } +} + +fn option_value<'a>(args: &'a [String], key: &str) -> Option<&'a str> { + args.windows(2).find(|pair| pair[0] == key).map(|pair| pair[1].as_str()) +} +``` + +- [ ] **Step 3: Verify CLI compiles** + +Run: + +```bash +cargo check -p dbx-cli +``` + +Expected: CLI crate compiles. + +--- + +### Task 7: Headless Connection Commands + +**Files:** +- Modify: `/Users/bytedance/open_source_poj/dbx/crates/dbx-cli/src/commands.rs` + +- [ ] **Step 1: Add headless state helper and redaction** + +Add helper functions to `commands.rs`: + +```rust +async fn open_state() -> Result { + let db_path = crate::runtime_client::app_data_dir().join("dbx.db"); + let storage = dbx_core::storage::Storage::open(&db_path).await?; + Ok(dbx_core::connection::AppState::new(storage)) +} + +fn redacted_config(config: &dbx_core::models::connection::ConnectionConfig) -> serde_json::Value { + serde_json::json!({ + "id": config.id, + "name": config.name, + "databaseType": config.db_type, + "driverProfile": config.driver_profile, + "driverLabel": config.driver_label, + "defaultDatabase": config.database, + "color": config.color, + "sshEnabled": config.ssh_enabled, + "redactedUrl": config.redacted_connection_url(), + }) +} + +async fn find_connection(name: &str) -> Result> { + let state = open_state().await.map_err(|e| fail(CliSource::Headless, CliErrorCode::InternalError, e, false))?; + let configs = state.storage.load_connections().await.map_err(|e| fail(CliSource::Headless, CliErrorCode::InternalError, e, false))?; + let matches: Vec<_> = configs.into_iter().filter(|c| c.name == name || c.id == name).collect(); + match matches.len() { + 1 => Ok(matches.into_iter().next().unwrap()), + 0 => Err(fail(CliSource::Headless, CliErrorCode::ConnectionNotFound, "Connection not found", true)), + _ => Err(fail(CliSource::Headless, CliErrorCode::AmbiguousConnection, "Connection name is ambiguous", true)), + } +} +``` + +- [ ] **Step 2: Replace `conn_list` implementation** + +Use this implementation: + +```rust +async fn conn_list(_args: &[String]) -> CliEnvelope { + let state = match open_state().await { + Ok(state) => state, + Err(e) => return fail(CliSource::Headless, CliErrorCode::InternalError, e, false), + }; + match state.storage.load_connections().await { + Ok(configs) => ok(CliSource::Headless, serde_json::json!({ + "connections": configs.iter().map(redacted_config).collect::>() + })), + Err(e) => fail(CliSource::Headless, CliErrorCode::InternalError, e, false), + } +} +``` + +- [ ] **Step 3: Replace `conn_show` implementation** + +Use this implementation: + +```rust +async fn conn_show(name: &str, _args: &[String]) -> CliEnvelope { + match find_connection(name).await { + Ok(config) => ok(CliSource::Headless, redacted_config(&config)), + Err(err) => err, + } +} +``` + +- [ ] **Step 4: Verify commands** + +Run: + +```bash +cargo run -p dbx-cli -- conn list --format json +cargo run -p dbx-cli -- conn show __missing__ --redacted --format json +``` + +Expected: first command returns an envelope; second returns `CONNECTION_NOT_FOUND` without panicking. + +--- + +### Task 8: Headless Schema Snapshot and Safe Query + +**Files:** +- Modify: `/Users/bytedance/open_source_poj/dbx/crates/dbx-cli/src/commands.rs` + +- [ ] **Step 1: Add connection registration helper** + +Add this helper: + +```rust +async fn state_with_connection(config: dbx_core::models::connection::ConnectionConfig) -> Result { + let state = open_state().await?; + state.configs.lock().await.insert(config.id.clone(), config.clone()); + state.get_or_create_pool(&config.id, config.database.as_deref()).await?; + Ok(state) +} +``` + +- [ ] **Step 2: Replace `schema_snapshot`** + +Use this implementation: + +```rust +async fn schema_snapshot(args: &[String]) -> CliEnvelope { + let Some(conn_name) = option_value(args, "--conn") else { + return fail(CliSource::Headless, CliErrorCode::ConnectionNotFound, "--conn is required", true); + }; + let db = option_value(args, "--db"); + let config = match find_connection(conn_name).await { + Ok(config) => config, + Err(err) => return err, + }; + let state = match state_with_connection(config.clone()).await { + Ok(state) => state, + Err(e) => return fail(CliSource::Headless, CliErrorCode::InternalError, e, false), + }; + match dbx_core::schema_snapshot::snapshot(&state, &config.id, db, None).await { + Ok(snapshot) => ok(CliSource::Headless, serde_json::to_value(snapshot).unwrap()), + Err(e) => fail(CliSource::Headless, CliErrorCode::InternalError, e, false), + } +} +``` + +- [ ] **Step 3: Replace `safe_query`** + +Use this implementation: + +```rust +async fn safe_query(args: &[String]) -> CliEnvelope { + let Some(conn_name) = option_value(args, "--conn") else { + return fail(CliSource::Headless, CliErrorCode::ConnectionNotFound, "--conn is required", true); + }; + let Some(sql) = option_value(args, "--sql") else { + return fail(CliSource::Headless, CliErrorCode::QueryClassificationFailed, "--sql is required", true); + }; + let config = match find_connection(conn_name).await { + Ok(config) => config, + Err(err) => return err, + }; + let risk = dbx_core::sql_safety::risk_for(sql, &config.name, config.color.as_deref()); + match risk.operation_class { + dbx_core::sql_safety::OperationClass::Read => {} + dbx_core::sql_safety::OperationClass::Write if risk.is_production => { + return fail(CliSource::Headless, CliErrorCode::ProductionWriteBlocked, serde_json::to_string(&risk).unwrap(), true); + } + dbx_core::sql_safety::OperationClass::Write => { + return fail(CliSource::Headless, CliErrorCode::HandoffRequired, serde_json::to_string(&risk).unwrap(), true); + } + dbx_core::sql_safety::OperationClass::Ddl => { + return fail(CliSource::Headless, CliErrorCode::DdlBlocked, serde_json::to_string(&risk).unwrap(), true); + } + dbx_core::sql_safety::OperationClass::Unknown => { + return fail(CliSource::Headless, CliErrorCode::QueryClassificationFailed, serde_json::to_string(&risk).unwrap(), true); + } + } + let state = match state_with_connection(config.clone()).await { + Ok(state) => state, + Err(e) => return fail(CliSource::Headless, CliErrorCode::InternalError, e, false), + }; + let database = option_value(args, "--db").or(config.database.as_deref()).unwrap_or(""); + match dbx_core::query::execute_sql_statement(&state, &config.id, database, sql, None, None).await { + Ok(result) => ok(CliSource::Headless, serde_json::json!({ "risk": risk, "result": result })), + Err(e) => fail(CliSource::Headless, CliErrorCode::InternalError, e, false), + } +} +``` + +- [ ] **Step 4: Verify SQL blocking** + +Run: + +```bash +cargo run -p dbx-cli -- safe-query --conn __missing__ --sql "DROP TABLE users" --format json +``` + +Expected: returns `CONNECTION_NOT_FOUND`. With a real DBX connection, `DROP TABLE users` returns `DDL_BLOCKED`. + +--- + +### Task 9: GUI Runtime Server + +**Files:** +- Create: `/Users/bytedance/open_source_poj/dbx/src-tauri/src/commands/agent_runtime.rs` +- Modify: `/Users/bytedance/open_source_poj/dbx/src-tauri/src/commands/mod.rs` +- Modify: `/Users/bytedance/open_source_poj/dbx/src-tauri/src/lib.rs` + +- [ ] **Step 1: Add runtime command module** + +Create `/Users/bytedance/open_source_poj/dbx/src-tauri/src/commands/agent_runtime.rs` with: + +```rust +use serde::{Deserialize, Serialize}; +use std::sync::Arc; +use tauri::{AppHandle, Manager}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::TcpListener; +use tokio::sync::RwLock; + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AgentRuntimeSnapshot { + pub active_connection_id: Option, + pub active_connection_name: Option, + pub database: Option, + pub schema: Option, + pub active_tab_id: Option, + pub active_tab_title: Option, + pub sql: Option, + pub selected_sql: Option, + pub selection: Option, + pub result: Option, +} + +#[derive(Clone)] +pub struct AgentRuntimeState { + pub token: String, + pub snapshot: Arc>, + pub handoffs: Arc>>, +} + +#[tauri::command] +pub async fn agent_runtime_update_snapshot( + state: tauri::State<'_, AgentRuntimeState>, + snapshot: AgentRuntimeSnapshot, +) -> Result<(), String> { + *state.snapshot.write().await = snapshot; + Ok(()) +} + +#[tauri::command] +pub async fn agent_runtime_load_handoffs( + app_state: tauri::State<'_, Arc>, + runtime: tauri::State<'_, AgentRuntimeState>, +) -> Result, String> { + let mut items = app_state.storage.load_pending_handoffs().await?; + items.extend(runtime.handoffs.read().await.iter().cloned()); + Ok(items) +} + +pub fn start(app: AppHandle) -> AgentRuntimeState { + let token = uuid::Uuid::new_v4().to_string(); + let state = AgentRuntimeState { + token: token.clone(), + snapshot: Arc::new(RwLock::new(AgentRuntimeSnapshot::default())), + handoffs: Arc::new(RwLock::new(Vec::new())), + }; + let server_state = state.clone(); + tauri::async_runtime::spawn(async move { + let listener = match TcpListener::bind("127.0.0.1:0").await { + Ok(listener) => listener, + Err(err) => { + log::warn!("Agent runtime bind failed: {err}"); + return; + } + }; + let port = listener.local_addr().map(|addr| addr.port()).unwrap_or(0); + if let Ok(dir) = app.path().app_data_dir() { + let payload = serde_json::json!({ "port": port, "token": token }); + let _ = std::fs::write(dir.join("agent-runtime.json"), serde_json::to_string(&payload).unwrap()); + } + loop { + let Ok((stream, _)) = listener.accept().await else { continue }; + let st = server_state.clone(); + tauri::async_runtime::spawn(async move { + handle_connection(stream, st).await; + }); + } + }); + state +} + +async fn handle_connection(mut stream: tokio::net::TcpStream, state: AgentRuntimeState) { + let mut buf = vec![0u8; 65536]; + let Ok(n) = stream.read(&mut buf).await else { return }; + if n == 0 { + return; + } + let request = String::from_utf8_lossy(&buf[..n]); + if !request.contains(&format!("Authorization: Bearer {}", state.token)) { + respond_json(&mut stream, "401 Unauthorized", serde_json::json!({"error":"unauthorized"})).await; + return; + } + let first = request.lines().next().unwrap_or(""); + if first.starts_with("GET /context") { + respond_json(&mut stream, "200 OK", serde_json::to_value(&*state.snapshot.read().await).unwrap()).await; + } else if first.starts_with("GET /selection") { + let snapshot = state.snapshot.read().await; + respond_json(&mut stream, "200 OK", snapshot.selection.clone().unwrap_or_else(|| serde_json::json!({"type":"none"}))).await; + } else if first.starts_with("GET /result/current") { + let snapshot = state.snapshot.read().await; + respond_json(&mut stream, "200 OK", snapshot.result.clone().unwrap_or_else(|| serde_json::json!({"columns":[],"rows":[]}))).await; + } else if first.starts_with("POST /handoff") { + let body = request.split("\r\n\r\n").nth(1).unwrap_or(""); + if let Ok(item) = serde_json::from_str::(body) { + state.handoffs.write().await.push(item.clone()); + respond_json(&mut stream, "200 OK", serde_json::json!({"id": item.id, "status": "shown"})).await; + } else { + respond_json(&mut stream, "400 Bad Request", serde_json::json!({"error":"invalid handoff"})).await; + } + } else { + respond_json(&mut stream, "404 Not Found", serde_json::json!({"error":"not found"})).await; + } +} + +async fn respond_json(stream: &mut tokio::net::TcpStream, status: &str, body: serde_json::Value) { + let body = serde_json::to_string(&body).unwrap(); + let resp = format!("HTTP/1.1 {status}\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{body}", body.len()); + let _ = stream.write_all(resp.as_bytes()).await; +} +``` + +- [ ] **Step 2: Expose module** + +Add to `/Users/bytedance/open_source_poj/dbx/src-tauri/src/commands/mod.rs`: + +```rust +pub mod agent_runtime; +``` + +- [ ] **Step 3: Start runtime in Tauri setup** + +In `/Users/bytedance/open_source_poj/dbx/src-tauri/src/lib.rs`, after `commands::mcp_bridge::start(app_handle, state);`, add: + +```rust +let runtime_state = commands::agent_runtime::start(app.handle().clone()); +app.manage(runtime_state); +``` + +Add command to `invoke_handler`: + +```rust +commands::agent_runtime::agent_runtime_update_snapshot, +commands::agent_runtime::agent_runtime_load_handoffs, +``` + +- [ ] **Step 4: Compile Tauri crate** + +Run: + +```bash +cargo check -p dbx +``` + +Expected: desktop crate compiles. + +--- + +### Task 10: Frontend Runtime State Sync + +**Files:** +- Modify: `/Users/bytedance/open_source_poj/dbx/src/lib/tauri.ts` +- Modify: `/Users/bytedance/open_source_poj/dbx/src/lib/api.ts` +- Create: `/Users/bytedance/open_source_poj/dbx/src/stores/agentRuntimeStore.ts` +- Modify: `/Users/bytedance/open_source_poj/dbx/src/stores/queryStore.ts` +- Modify: `/Users/bytedance/open_source_poj/dbx/src/components/grid/DataGrid.vue` + +- [ ] **Step 1: Add Tauri wrappers** + +Add to `/Users/bytedance/open_source_poj/dbx/src/lib/tauri.ts`: + +```ts +export async function agentRuntimeUpdateSnapshot(snapshot: unknown): Promise { + return invoke("agent_runtime_update_snapshot", { snapshot }); +} + +export async function agentRuntimeLoadHandoffs(): Promise { + return invoke("agent_runtime_load_handoffs"); +} +``` + +- [ ] **Step 2: Export API forwards** + +Add to `/Users/bytedance/open_source_poj/dbx/src/lib/api.ts`: + +```ts +export const agentRuntimeUpdateSnapshot = forward("agentRuntimeUpdateSnapshot"); +export const agentRuntimeLoadHandoffs = forward("agentRuntimeLoadHandoffs"); +``` + +- [ ] **Step 3: Create runtime store** + +Create `/Users/bytedance/open_source_poj/dbx/src/stores/agentRuntimeStore.ts`: + +```ts +import { defineStore } from "pinia"; +import { ref } from "vue"; +import * as api from "@/lib/api"; +import { useConnectionStore } from "@/stores/connectionStore"; +import { useQueryStore } from "@/stores/queryStore"; + +export const useAgentRuntimeStore = defineStore("agentRuntime", () => { + const selection = ref({ type: "none" }); + let timer: ReturnType | null = null; + + function setSelection(value: unknown) { + selection.value = value; + scheduleSync(); + } + + function scheduleSync() { + if (timer) clearTimeout(timer); + timer = setTimeout(() => void syncNow(), 100); + } + + async function syncNow() { + const connectionStore = useConnectionStore(); + const queryStore = useQueryStore(); + const tab = queryStore.tabs.find((item) => item.id === queryStore.activeTabId); + const conn = tab ? connectionStore.getConfig(tab.connectionId) : undefined; + await api.agentRuntimeUpdateSnapshot({ + activeConnectionId: tab?.connectionId, + activeConnectionName: conn?.name, + database: tab?.database, + schema: tab?.schema, + activeTabId: tab?.id, + activeTabTitle: tab?.title, + sql: tab?.sql, + selectedSql: undefined, + selection: selection.value, + result: tab?.result + ? { + columns: tab.result.columns, + rows: tab.result.rows.slice(0, 50), + truncated: tab.result.rows.length > 50 || tab.result.truncated, + executionTimeMs: tab.result.execution_time_ms, + } + : undefined, + }); + } + + return { selection, setSelection, scheduleSync, syncNow }; +}); +``` + +- [ ] **Step 4: Notify sync from query store** + +In `/Users/bytedance/open_source_poj/dbx/src/stores/queryStore.ts`, import: + +```ts +import { useAgentRuntimeStore } from "@/stores/agentRuntimeStore"; +``` + +After result assignment in `executeTabSql`, call: + +```ts +useAgentRuntimeStore().scheduleSync(); +``` + +Also call `useAgentRuntimeStore().scheduleSync();` after active tab changes in `createTab`, `openSavedSql`, `closeTab`, and `setActiveResultIndex`. + +- [ ] **Step 5: Notify grid selection** + +In `/Users/bytedance/open_source_poj/dbx/src/components/grid/DataGrid.vue`, import: + +```ts +import { useAgentRuntimeStore } from "@/stores/agentRuntimeStore"; +``` + +Create a store instance: + +```ts +const agentRuntimeStore = useAgentRuntimeStore(); +``` + +Where selection range changes, call: + +```ts +agentRuntimeStore.setSelection({ + type: "grid-cells", + data: extractSelection(props.result.columns, props.result.rows, selectionRange.value), +}); +``` + +- [ ] **Step 6: Run frontend typecheck** + +Run: + +```bash +pnpm build +``` + +Expected: Vue typecheck and Vite build pass. + +--- + +### Task 11: Handoff CLI and GUI Review UI + +**Files:** +- Modify: `/Users/bytedance/open_source_poj/dbx/crates/dbx-cli/src/commands.rs` +- Create: `/Users/bytedance/open_source_poj/dbx/src/components/agent/AgentHandoffDialog.vue` +- Modify: `/Users/bytedance/open_source_poj/dbx/src/components/layout/AppDialogs.vue` + +- [ ] **Step 1: Implement CLI handoff** + +Replace `handoff` in `commands.rs` with: + +```rust +async fn handoff(args: &[String]) -> CliEnvelope { + let Some(conn_name) = option_value(args, "--conn") else { + return fail(CliSource::Headless, CliErrorCode::ConnectionNotFound, "--conn is required", true); + }; + let Some(title) = option_value(args, "--title") else { + return fail(CliSource::Headless, CliErrorCode::InternalError, "--title is required", true); + }; + let sql = if let Some(sql_file) = option_value(args, "--sql-file") { + match std::fs::read_to_string(sql_file) { + Ok(sql) => sql, + Err(e) => return fail(CliSource::Headless, CliErrorCode::InternalError, e.to_string(), true), + } + } else if let Some(sql) = option_value(args, "--sql") { + sql.to_string() + } else { + return fail(CliSource::Headless, CliErrorCode::InternalError, "--sql-file or --sql is required", true); + }; + let config = match find_connection(conn_name).await { + Ok(config) => config, + Err(err) => return err, + }; + let risk = dbx_core::sql_safety::risk_for(&sql, &config.name, config.color.as_deref()); + let item = dbx_core::handoff::HandoffItem::queued( + config.name, + config.database, + title.to_string(), + option_value(args, "--description").map(str::to_string), + sql, + risk.operation_class, + risk.risk_level, + risk.is_production, + ); + if let Ok(data) = crate::runtime_client::post_json("/handoff", serde_json::to_value(&item).unwrap()).await { + return ok(CliSource::GuiRuntime, data); + } + match open_state().await.and_then(|state| futures::executor::block_on(state.storage.save_handoff(&item))) { + Ok(()) => ok(CliSource::Headless, serde_json::json!({ "id": item.id, "status": "queued" })), + Err(e) => fail(CliSource::Headless, CliErrorCode::InternalError, e, false), + } +} +``` + +- [ ] **Step 2: Create handoff review dialog** + +Create `/Users/bytedance/open_source_poj/dbx/src/components/agent/AgentHandoffDialog.vue`: + +```vue + + + +``` + +- [ ] **Step 3: Mount dialog** + +In `/Users/bytedance/open_source_poj/dbx/src/components/layout/AppDialogs.vue`, import and add: + +```vue + +``` + +- [ ] **Step 4: Verify GUI build** + +Run: + +```bash +pnpm build +``` + +Expected: build passes and dialog component compiles. + +--- + +### Task 12: Runtime Context, Selection, and Result CLI Verification + +**Files:** +- Modify: `/Users/bytedance/open_source_poj/dbx/crates/dbx-cli/src/commands.rs` +- Modify: `/Users/bytedance/open_source_poj/dbx/src/stores/agentRuntimeStore.ts` + +- [ ] **Step 1: Improve context headless fallback** + +Replace headless branch of `context` with: + +```rust +Err(_) => { + let state = match open_state().await { + Ok(state) => state, + Err(e) => return fail(CliSource::Headless, CliErrorCode::InternalError, e, false), + }; + let configs = state.storage.load_connections().await.unwrap_or_default(); + ok(CliSource::Headless, serde_json::json!({ + "runtime": "headless", + "activeConnection": configs.first().map(redacted_config), + "configSource": crate::runtime_client::app_data_dir().join("dbx.db").display().to_string() + })) +} +``` + +- [ ] **Step 2: Ensure result limit is applied** + +In `/Users/bytedance/open_source_poj/dbx/src/stores/agentRuntimeStore.ts`, keep `rows.slice(0, 50)` and ensure `truncated` is set when original rows exceed 50. + +- [ ] **Step 3: Manual runtime verification** + +Run DBX desktop: + +```bash +pnpm tauri dev +``` + +In a second terminal, run: + +```bash +cargo run -p dbx-cli -- context --format json +cargo run -p dbx-cli -- selection --format json +cargo run -p dbx-cli -- result current --limit 50 --format json +``` + +Expected: all three commands return `source: "gui-runtime"` when desktop is running. + +--- + +### Task 13: Final Verification and Compatibility Check + +**Files:** +- No new files. + +- [ ] **Step 1: Verify Rust workspace** + +Run: + +```bash +cargo check --workspace +cargo test -p dbx-core +``` + +Expected: workspace checks and core tests pass. + +- [ ] **Step 2: Verify frontend** + +Run: + +```bash +pnpm build +``` + +Expected: Vue typecheck and Vite build pass. + +- [ ] **Step 3: Verify MCP untouched** + +Run: + +```bash +git diff -- mcp +``` + +Expected: no output. + +- [ ] **Step 4: Verify CLI command surface** + +Run: + +```bash +cargo run -p dbx-cli -- context --format json +cargo run -p dbx-cli -- conn list --format json +cargo run -p dbx-cli -- conn show __missing__ --redacted --format json +cargo run -p dbx-cli -- selection --format json +cargo run -p dbx-cli -- result current --limit 50 --format json +``` + +Expected: all commands return valid JSON envelopes. GUI-only commands return `GUI_RUNTIME_REQUIRED` if DBX desktop is not running. + +## Self-Review + +- Spec coverage: The plan covers all 8 CLI commands, runtime discovery/token, GUI context, selection, current result, safe-query classification, schema snapshot, and handoff queue/display. +- Compatibility: The plan explicitly avoids changing `mcp/` and verifies this with `git diff -- mcp`. +- Type consistency: Rust DTOs use camelCase where sent to frontend/CLI. Error codes use screaming snake case. Source uses kebab case. +- Known implementation detail to watch: Task 11 uses a synchronous block inside async CLI code for queued handoff save; if it causes compile friction, replace it with direct async `state.storage.save_handoff(&item).await` inside a `match` after opening state. From 5720f05dd08581573e53f0d961345a1d6c270737 Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 15:47:52 +0800 Subject: [PATCH 03/40] =?UTF-8?q?feat(cli):=20=E6=B7=BB=E5=8A=A0=20CLI=20c?= =?UTF-8?q?rate=20=E9=AA=A8=E6=9E=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- Cargo.toml | 2 +- crates/dbx-cli/Cargo.toml | 17 +++++++++++++++++ crates/dbx-cli/src/main.rs | 13 +++++++++++++ 3 files changed, 31 insertions(+), 1 deletion(-) create mode 100644 crates/dbx-cli/Cargo.toml create mode 100644 crates/dbx-cli/src/main.rs diff --git a/Cargo.toml b/Cargo.toml index 390ef8ef7..1b5c694ec 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [workspace] resolver = "2" -members = ["src-tauri", "crates/dbx-core", "src-web"] +members = ["src-tauri", "crates/dbx-core", "crates/dbx-cli", "src-web"] [profile.release] panic = "abort" diff --git a/crates/dbx-cli/Cargo.toml b/crates/dbx-cli/Cargo.toml new file mode 100644 index 000000000..c816811ec --- /dev/null +++ b/crates/dbx-cli/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "dbx-cli" +version = "0.5.2" +edition = "2021" + +[[bin]] +name = "dbx" +path = "src/main.rs" + +[dependencies] +dbx-core = { path = "../dbx-core" } +chrono = { version = "0.4", features = ["serde"] } +serde = { version = "1.0", features = ["derive"] } +serde_json = "1.0" +tokio = { version = "1", features = ["full"] } +reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] } +uuid = { version = "1", features = ["v4", "serde"] } diff --git a/crates/dbx-cli/src/main.rs b/crates/dbx-cli/src/main.rs new file mode 100644 index 000000000..f34308a75 --- /dev/null +++ b/crates/dbx-cli/src/main.rs @@ -0,0 +1,13 @@ +mod commands; +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()) + ); + std::process::exit(1); + } +} From 01110945a1cb7a6f0d94007b56e69c8d459bd091 Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 15:56:44 +0800 Subject: [PATCH 04/40] =?UTF-8?q?fix(dbx-cli):=20=E8=A1=A5=E9=BD=90?= =?UTF-8?q?=E6=9C=80=E5=B0=8F=20runtime=20client=20=E6=A8=A1=E5=9D=97?= =?UTF-8?q?=E4=BB=A5=E6=81=A2=E5=A4=8D=E7=BC=96=E8=AF=91?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 新增 commands 和 runtime_client 最小实现，补足 main.rs 所需模块 - 允许 context 读取 GUI runtime，GUI 专属命令在 headless 下返回结构化错误 - 更新 Cargo.lock 并确认 cargo check -p dbx-cli --locked 通过 --- Cargo.lock | 13 ++++++ crates/dbx-cli/src/commands.rs | 61 ++++++++++++++++++++++++++++ crates/dbx-cli/src/runtime_client.rs | 42 +++++++++++++++++++ 3 files changed, 116 insertions(+) create mode 100644 crates/dbx-cli/src/commands.rs create mode 100644 crates/dbx-cli/src/runtime_client.rs diff --git a/Cargo.lock b/Cargo.lock index c8f06c8e7..27dfd5fef 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1673,6 +1673,19 @@ dependencies = [ "zip 4.6.1", ] +[[package]] +name = "dbx-cli" +version = "0.5.2" +dependencies = [ + "chrono", + "dbx-core", + "reqwest 0.12.28", + "serde", + "serde_json", + "tokio", + "uuid", +] + [[package]] name = "dbx-core" version = "0.1.0" diff --git a/crates/dbx-cli/src/commands.rs b/crates/dbx-cli/src/commands.rs new file mode 100644 index 000000000..60d828322 --- /dev/null +++ b/crates/dbx-cli/src/commands.rs @@ -0,0 +1,61 @@ +use serde::Serialize; + +#[derive(Debug, Serialize)] +#[serde(rename_all = "SCREAMING_SNAKE_CASE")] +pub(crate) enum CliErrorCode { + GuiRuntimeRequired, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "kebab-case")] +pub(crate) enum CliSource { + GuiRuntime, + Headless, +} + +#[derive(Debug, Serialize)] +pub(crate) struct CliError { + code: CliErrorCode, + message: String, + recoverable: bool, +} + +#[derive(Debug, Serialize)] +#[serde(untagged)] +pub(crate) enum CliEnvelope { + Success { ok: bool, source: CliSource, data: T }, + Failure { ok: bool, source: CliSource, error: CliError }, +} + +pub(crate) async fn run(args: Vec) -> Result<(), CliEnvelope<()>> { + let output = match args.as_slice() { + [cmd, ..] if cmd == "context" => context().await, + [cmd, ..] if cmd == "selection" => runtime_required("dbx selection requires DBX GUI runtime."), + [cmd, sub, ..] if cmd == "result" && sub == "current" => { + runtime_required("dbx result current requires DBX GUI runtime.") + } + _ => ok(CliSource::Headless, serde_json::json!({ "runtime": "headless" })), + }; + + println!("{}", serde_json::to_string_pretty(&output).unwrap()); + Ok(()) +} + +async fn context() -> CliEnvelope { + match crate::runtime_client::get_json("/context").await { + Ok(data) => ok(CliSource::GuiRuntime, data), + Err(_) => ok(CliSource::Headless, serde_json::json!({ "runtime": "headless" })), + } +} + +fn ok(source: CliSource, data: T) -> CliEnvelope { + CliEnvelope::Success { ok: true, source, data } +} + +fn runtime_required(message: &str) -> CliEnvelope { + CliEnvelope::Failure { + ok: false, + source: CliSource::Headless, + error: CliError { code: CliErrorCode::GuiRuntimeRequired, message: message.to_string(), recoverable: true }, + } +} diff --git a/crates/dbx-cli/src/runtime_client.rs b/crates/dbx-cli/src/runtime_client.rs new file mode 100644 index 000000000..a3db99811 --- /dev/null +++ b/crates/dbx-cli/src/runtime_client.rs @@ -0,0 +1,42 @@ +use serde::{Deserialize, Serialize}; +use std::path::PathBuf; + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct RuntimeDiscovery { + pub port: u16, + pub token: String, +} + +pub fn app_data_dir() -> PathBuf { + 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") + } else if cfg!(windows) { + PathBuf::from(home).join("com.dbx.app") + } else { + PathBuf::from(home).join(".config/com.dbx.app") + } +} + +pub fn load_runtime() -> Option { + let path = app_data_dir().join("agent-runtime.json"); + let json = std::fs::read_to_string(path).ok()?; + serde_json::from_str(&json).ok() +} + +pub async fn get_json(path: &str) -> Result { + let runtime = load_runtime().ok_or_else(|| "runtime unavailable".to_string())?; + let url = format!("http://127.0.0.1:{}{}", runtime.port, path); + + reqwest::Client::new() + .get(url) + .bearer_auth(runtime.token) + .send() + .await + .map_err(|err| err.to_string())? + .json() + .await + .map_err(|err| err.to_string()) +} From dd113ed7146c4a839853bbb0101d0529b0f16669 Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 16:45:04 +0800 Subject: [PATCH 05/40] =?UTF-8?q?build(dbx-cli):=20=E5=B0=86=20CLI=20?= =?UTF-8?q?=E4=BA=8C=E8=BF=9B=E5=88=B6=E6=94=B9=E5=90=8D=E4=B8=BA=20dbx-cl?= =?UTF-8?q?i?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 更新 dbx-cli crate 的 bin 目标名，避免与桌面 dbx 目标冲突 - 同步计划文档中的 cargo run 命令引用 - 保留 runtime_client.rs 现有格式化改动 --- crates/dbx-cli/Cargo.toml | 2 +- crates/dbx-cli/src/runtime_client.rs | 3 ++- .../plans/2026-05-10-dbx-cli-runtime.md | 26 +++++++++---------- 3 files changed, 16 insertions(+), 15 deletions(-) diff --git a/crates/dbx-cli/Cargo.toml b/crates/dbx-cli/Cargo.toml index c816811ec..d272aaa4b 100644 --- a/crates/dbx-cli/Cargo.toml +++ b/crates/dbx-cli/Cargo.toml @@ -4,7 +4,7 @@ version = "0.5.2" edition = "2021" [[bin]] -name = "dbx" +name = "dbx-cli" path = "src/main.rs" [dependencies] diff --git a/crates/dbx-cli/src/runtime_client.rs b/crates/dbx-cli/src/runtime_client.rs index a3db99811..940fd5399 100644 --- a/crates/dbx-cli/src/runtime_client.rs +++ b/crates/dbx-cli/src/runtime_client.rs @@ -9,7 +9,8 @@ pub struct RuntimeDiscovery { } pub fn app_data_dir() -> PathBuf { - 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") diff --git a/docs/superpowers/plans/2026-05-10-dbx-cli-runtime.md b/docs/superpowers/plans/2026-05-10-dbx-cli-runtime.md index 961f99829..e3cd3baf7 100644 --- a/docs/superpowers/plans/2026-05-10-dbx-cli-runtime.md +++ b/docs/superpowers/plans/2026-05-10-dbx-cli-runtime.md @@ -76,7 +76,7 @@ version = "0.5.2" edition = "2021" [[bin]] -name = "dbx" +name = "dbx-cli" path = "src/main.rs" [dependencies] @@ -114,7 +114,7 @@ Run: cargo metadata --format-version 1 --no-deps ``` -Expected: output contains `"name":"dbx-cli"` and `"name":"dbx"`. +Expected: output contains `"name":"dbx-cli"` for the package and `"name":"dbx-cli"` for the binary target. --- @@ -874,8 +874,8 @@ async fn conn_show(name: &str, _args: &[String]) -> CliEnvelope CliEnvelope { Run: ```bash -cargo run -p dbx-cli -- safe-query --conn __missing__ --sql "DROP TABLE users" --format json +cargo run -p dbx-cli --bin dbx-cli -- safe-query --conn __missing__ --sql "DROP TABLE users" --format json ``` Expected: returns `CONNECTION_NOT_FOUND`. With a real DBX connection, `DROP TABLE users` returns `DDL_BLOCKED`. @@ -1470,9 +1470,9 @@ pnpm tauri dev In a second terminal, run: ```bash -cargo run -p dbx-cli -- context --format json -cargo run -p dbx-cli -- selection --format json -cargo run -p dbx-cli -- result current --limit 50 --format json +cargo run -p dbx-cli --bin dbx-cli -- context --format json +cargo run -p dbx-cli --bin dbx-cli -- selection --format json +cargo run -p dbx-cli --bin dbx-cli -- result current --limit 50 --format json ``` Expected: all three commands return `source: "gui-runtime"` when desktop is running. @@ -1520,11 +1520,11 @@ Expected: no output. Run: ```bash -cargo run -p dbx-cli -- context --format json -cargo run -p dbx-cli -- conn list --format json -cargo run -p dbx-cli -- conn show __missing__ --redacted --format json -cargo run -p dbx-cli -- selection --format json -cargo run -p dbx-cli -- result current --limit 50 --format json +cargo run -p dbx-cli --bin dbx-cli -- context --format json +cargo run -p dbx-cli --bin dbx-cli -- conn list --format json +cargo run -p dbx-cli --bin dbx-cli -- conn show __missing__ --redacted --format json +cargo run -p dbx-cli --bin dbx-cli -- selection --format json +cargo run -p dbx-cli --bin dbx-cli -- result current --limit 50 --format json ``` Expected: all commands return valid JSON envelopes. GUI-only commands return `GUI_RUNTIME_REQUIRED` if DBX desktop is not running. From 10373bdf36a9f879199a4756f62ccbec0fc89a69 Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 16:52:39 +0800 Subject: [PATCH 06/40] =?UTF-8?q?feat(dbx-core):=20=E6=96=B0=E5=A2=9E=20CL?= =?UTF-8?q?I=20envelope=20DTO?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 添加 CLI source、error code、error 和 envelope 类型 - 导出 dbx-core::cli 模块并补充序列化单元测试 - 补齐既有测试 fixture 的 JDBC 默认字段以恢复目标测试编译 --- crates/dbx-core/src/cli.rs | 67 ++++++++++++++++++++++++++++++++++++++ crates/dbx-core/src/lib.rs | 1 + 2 files changed, 68 insertions(+) create mode 100644 crates/dbx-core/src/cli.rs diff --git a/crates/dbx-core/src/cli.rs b/crates/dbx-core/src/cli.rs new file mode 100644 index 000000000..f57932f6c --- /dev/null +++ b/crates/dbx-core/src/cli.rs @@ -0,0 +1,67 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "kebab-case")] +pub enum CliSource { + GuiRuntime, + Headless, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "SCREAMING_SNAKE_CASE")] +pub enum CliErrorCode { + GuiRuntimeRequired, + ConnectionNotFound, + AmbiguousConnection, + SecretUnavailable, + SshTunnelFailed, + QueryClassificationFailed, + HandoffRequired, + DdlBlocked, + ProductionWriteBlocked, + UnsupportedDatabaseType, + Timeout, + InternalError, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CliError { + pub code: CliErrorCode, + pub message: String, + pub recoverable: bool, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(untagged)] +pub enum CliEnvelope { + Success { ok: bool, source: CliSource, data: T }, + Failure { ok: bool, source: CliSource, error: CliError }, +} + +pub fn ok(source: CliSource, data: T) -> CliEnvelope { + CliEnvelope::Success { ok: true, source, data } +} + +pub fn fail(source: CliSource, code: CliErrorCode, message: impl Into, recoverable: bool) -> CliEnvelope { + CliEnvelope::Failure { ok: false, source, error: CliError { code, message: message.into(), recoverable } } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn serializes_success_source_as_kebab_case() { + let env = ok(CliSource::GuiRuntime, serde_json::json!({"value": 1})); + let json = serde_json::to_string(&env).unwrap(); + assert!(json.contains("\"ok\":true")); + assert!(json.contains("\"source\":\"gui-runtime\"")); + } + + #[test] + fn serializes_error_code_as_screaming_snake_case() { + let env: CliEnvelope<()> = fail(CliSource::Headless, CliErrorCode::GuiRuntimeRequired, "runtime needed", true); + let json = serde_json::to_string(&env).unwrap(); + assert!(json.contains("\"GUI_RUNTIME_REQUIRED\"")); + } +} diff --git a/crates/dbx-core/src/lib.rs b/crates/dbx-core/src/lib.rs index cd34fdf92..65ae7335a 100644 --- a/crates/dbx-core/src/lib.rs +++ b/crates/dbx-core/src/lib.rs @@ -1,4 +1,5 @@ pub mod ai; +pub mod cli; pub mod connection; pub mod connection_secrets; pub mod db; From 7f837390405c5662fb05da91878a78402a34189a Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 16:57:54 +0800 Subject: [PATCH 07/40] =?UTF-8?q?fix(cli):=20=E6=A0=A1=E9=AA=8C=20CLI=20en?= =?UTF-8?q?velope=20=E5=8F=8D=E5=BA=8F=E5=88=97=E5=8C=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 为 CliEnvelope 增加 ok/data/error 一致性校验和回归测试 - 让 dbx-cli commands 复用 dbx-core 的 CLI DTO 与 helper --- crates/dbx-cli/src/commands.rs | 39 +----- crates/dbx-core/src/cli.rs | 223 ++++++++++++++++++++++++++++++++- 2 files changed, 223 insertions(+), 39 deletions(-) diff --git a/crates/dbx-cli/src/commands.rs b/crates/dbx-cli/src/commands.rs index 60d828322..9ab5ab8d0 100644 --- a/crates/dbx-cli/src/commands.rs +++ b/crates/dbx-cli/src/commands.rs @@ -1,31 +1,4 @@ -use serde::Serialize; - -#[derive(Debug, Serialize)] -#[serde(rename_all = "SCREAMING_SNAKE_CASE")] -pub(crate) enum CliErrorCode { - GuiRuntimeRequired, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "kebab-case")] -pub(crate) enum CliSource { - GuiRuntime, - Headless, -} - -#[derive(Debug, Serialize)] -pub(crate) struct CliError { - code: CliErrorCode, - message: String, - recoverable: bool, -} - -#[derive(Debug, Serialize)] -#[serde(untagged)] -pub(crate) enum CliEnvelope { - Success { ok: bool, source: CliSource, data: T }, - Failure { ok: bool, source: CliSource, error: CliError }, -} +use dbx_core::cli::{fail, ok, CliEnvelope, CliErrorCode, CliSource}; pub(crate) async fn run(args: Vec) -> Result<(), CliEnvelope<()>> { let output = match args.as_slice() { @@ -48,14 +21,6 @@ async fn context() -> CliEnvelope { } } -fn ok(source: CliSource, data: T) -> CliEnvelope { - CliEnvelope::Success { ok: true, source, data } -} - fn runtime_required(message: &str) -> CliEnvelope { - CliEnvelope::Failure { - ok: false, - source: CliSource::Headless, - error: CliError { code: CliErrorCode::GuiRuntimeRequired, message: message.to_string(), recoverable: true }, - } + fail(CliSource::Headless, CliErrorCode::GuiRuntimeRequired, message, true) } diff --git a/crates/dbx-core/src/cli.rs b/crates/dbx-core/src/cli.rs index f57932f6c..26af58de2 100644 --- a/crates/dbx-core/src/cli.rs +++ b/crates/dbx-core/src/cli.rs @@ -1,4 +1,7 @@ -use serde::{Deserialize, Serialize}; +use serde::de::{self, IgnoredAny, MapAccess, Visitor}; +use serde::{Deserialize, Deserializer, Serialize}; +use std::fmt; +use std::marker::PhantomData; #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] #[serde(rename_all = "kebab-case")] @@ -31,13 +34,140 @@ pub struct CliError { pub recoverable: bool, } -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize)] #[serde(untagged)] pub enum CliEnvelope { Success { ok: bool, source: CliSource, data: T }, Failure { ok: bool, source: CliSource, error: CliError }, } +impl<'de, T> Deserialize<'de> for CliEnvelope +where + T: Deserialize<'de>, +{ + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + deserializer.deserialize_map(CliEnvelopeVisitor { marker: PhantomData }) + } +} + +struct CliEnvelopeVisitor { + marker: PhantomData, +} + +impl<'de, T> Visitor<'de> for CliEnvelopeVisitor +where + T: Deserialize<'de>, +{ + type Value = CliEnvelope; + + fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result { + formatter.write_str("a CLI envelope with consistent ok/data/error fields") + } + + fn visit_map(self, mut map: M) -> Result + where + M: MapAccess<'de>, + { + let mut ok = None; + let mut source = None; + let mut data = None; + let mut has_data = false; + let mut error = None; + let mut has_error = false; + + while let Some(key) = map.next_key::()? { + match key { + CliEnvelopeField::Ok => { + if ok.is_some() { + return Err(de::Error::duplicate_field("ok")); + } + ok = Some(map.next_value()?); + } + CliEnvelopeField::Source => { + if source.is_some() { + return Err(de::Error::duplicate_field("source")); + } + source = Some(map.next_value()?); + } + CliEnvelopeField::Data => { + if has_data { + return Err(de::Error::duplicate_field("data")); + } + has_data = true; + data = Some(map.next_value()?); + } + CliEnvelopeField::Error => { + if has_error { + return Err(de::Error::duplicate_field("error")); + } + has_error = true; + error = Some(map.next_value()?); + } + CliEnvelopeField::Ignore => { + let _ = map.next_value::()?; + } + } + } + + let ok = ok.ok_or_else(|| de::Error::missing_field("ok"))?; + 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, _, _) => Err(de::Error::custom("ok=true requires data without error")), + (false, _, _) => Err(de::Error::custom("ok=false requires error without data")), + } + } +} + +enum CliEnvelopeField { + Ok, + Source, + Data, + Error, + Ignore, +} + +impl<'de> Deserialize<'de> for CliEnvelopeField { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + deserializer.deserialize_identifier(CliEnvelopeFieldVisitor) + } +} + +struct CliEnvelopeFieldVisitor; + +impl<'de> Visitor<'de> for CliEnvelopeFieldVisitor { + type Value = CliEnvelopeField; + + fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result { + formatter.write_str("a CLI envelope field") + } + + fn visit_str(self, value: &str) -> Result + where + E: de::Error, + { + Ok(match value { + "ok" => CliEnvelopeField::Ok, + "source" => CliEnvelopeField::Source, + "data" => CliEnvelopeField::Data, + "error" => CliEnvelopeField::Error, + _ => CliEnvelopeField::Ignore, + }) + } +} + pub fn ok(source: CliSource, data: T) -> CliEnvelope { CliEnvelope::Success { ok: true, source, data } } @@ -64,4 +194,93 @@ mod tests { let json = serde_json::to_string(&env).unwrap(); assert!(json.contains("\"GUI_RUNTIME_REQUIRED\"")); } + + #[test] + fn deserializes_success_when_ok_true_and_data_present() { + let env: CliEnvelope = serde_json::from_value(serde_json::json!({ + "ok": true, + "source": "gui-runtime", + "data": { "value": 1 } + })) + .unwrap(); + + match env { + CliEnvelope::Success { ok, source, data } => { + assert!(ok); + assert_eq!(source, CliSource::GuiRuntime); + assert_eq!(data, serde_json::json!({ "value": 1 })); + } + CliEnvelope::Failure { .. } => panic!("expected success envelope"), + } + } + + #[test] + fn deserializes_failure_when_ok_false_and_error_present() { + let env: CliEnvelope = serde_json::from_value(serde_json::json!({ + "ok": false, + "source": "headless", + "error": { + "code": "GUI_RUNTIME_REQUIRED", + "message": "runtime needed", + "recoverable": true + } + })) + .unwrap(); + + match env { + CliEnvelope::Failure { ok, source, error } => { + assert!(!ok); + assert_eq!(source, CliSource::Headless); + assert_eq!(error.code, CliErrorCode::GuiRuntimeRequired); + assert_eq!(error.message, "runtime needed"); + assert!(error.recoverable); + } + CliEnvelope::Success { .. } => panic!("expected failure envelope"), + } + } + + #[test] + fn rejects_success_shape_when_ok_is_false() { + let err = serde_json::from_value::>(serde_json::json!({ + "ok": false, + "source": "headless", + "data": { "runtime": "headless" } + })) + .unwrap_err(); + + assert!(err.to_string().contains("ok=false requires error without data")); + } + + #[test] + fn rejects_failure_shape_when_ok_is_true() { + let err = serde_json::from_value::>(serde_json::json!({ + "ok": true, + "source": "headless", + "error": { + "code": "GUI_RUNTIME_REQUIRED", + "message": "runtime needed", + "recoverable": true + } + })) + .unwrap_err(); + + assert!(err.to_string().contains("ok=true requires data without error")); + } + + #[test] + fn rejects_envelope_with_both_data_and_error() { + let err = serde_json::from_value::>(serde_json::json!({ + "ok": true, + "source": "gui-runtime", + "data": { "value": 1 }, + "error": { + "code": "INTERNAL_ERROR", + "message": "unexpected", + "recoverable": false + } + })) + .unwrap_err(); + + assert!(err.to_string().contains("ok=true requires data without error")); + } } From bc6898cd22aa91fd67658515b6eef6b00013526b Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 17:02:28 +0800 Subject: [PATCH 08/40] =?UTF-8?q?feat(dbx-core):=20=E6=96=B0=E5=A2=9E=20SQ?= =?UTF-8?q?L=20=E5=AE=89=E5=85=A8=E9=A3=8E=E9=99=A9=E5=88=86=E7=B1=BB?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 添加 sql_safety 模块，区分 READ、WRITE、DDL 和 Unknown - 增加生产连接信号识别和风险元数据输出 - 导出模块并补充分类器单元测试 --- crates/dbx-core/src/lib.rs | 1 + crates/dbx-core/src/sql_safety.rs | 139 ++++++++++++++++++++++++++++++ 2 files changed, 140 insertions(+) create mode 100644 crates/dbx-core/src/sql_safety.rs diff --git a/crates/dbx-core/src/lib.rs b/crates/dbx-core/src/lib.rs index 65ae7335a..0636c7e43 100644 --- a/crates/dbx-core/src/lib.rs +++ b/crates/dbx-core/src/lib.rs @@ -14,6 +14,7 @@ pub mod redis_ops; pub mod saved_sql; pub mod schema; pub mod sql; +pub mod sql_safety; pub mod storage; pub mod table_import; pub mod transfer; diff --git a/crates/dbx-core/src/sql_safety.rs b/crates/dbx-core/src/sql_safety.rs new file mode 100644 index 000000000..5f24ada2d --- /dev/null +++ b/crates/dbx-core/src/sql_safety.rs @@ -0,0 +1,139 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "lowercase")] +pub enum OperationClass { + Read, + Write, + Ddl, + Unknown, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "lowercase")] +pub enum RiskLevel { + Low, + Medium, + High, + Critical, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct RiskMetadata { + pub operation_class: OperationClass, + pub risk_level: RiskLevel, + pub is_production: bool, + pub production_reason: Option, + pub first_token: Option, +} + +pub fn classify_sql(sql: &str) -> OperationClass { + let token = first_executable_token(sql).map(|s| s.to_ascii_uppercase()); + match token.as_deref() { + Some("SELECT" | "SHOW" | "DESCRIBE" | "EXPLAIN" | "WITH") => OperationClass::Read, + Some("INSERT" | "UPDATE" | "DELETE" | "MERGE" | "REPLACE") => OperationClass::Write, + Some("CREATE" | "ALTER" | "DROP" | "TRUNCATE" | "RENAME") => OperationClass::Ddl, + _ => OperationClass::Unknown, + } +} + +pub fn risk_for(sql: &str, connection_name: &str, color: Option<&str>) -> RiskMetadata { + let operation_class = classify_sql(sql); + let (is_production, production_reason) = production_signal(connection_name, color); + let risk_level = match (operation_class, is_production) { + (OperationClass::Read, _) => RiskLevel::Low, + (OperationClass::Write, false) => RiskLevel::Medium, + (OperationClass::Write, true) => RiskLevel::High, + (OperationClass::Ddl, _) => RiskLevel::Critical, + (OperationClass::Unknown, _) => RiskLevel::High, + }; + + RiskMetadata { + operation_class, + risk_level, + is_production, + production_reason, + first_token: first_executable_token(sql).map(str::to_string), + } +} + +fn production_signal(connection_name: &str, color: Option<&str>) -> (bool, Option) { + if matches!(color, Some("#ef4444")) { + return (true, Some("red connection color".to_string())); + } + + let name = connection_name.to_ascii_lowercase(); + if ["prod", "production", "live"] + .iter() + .any(|needle| name.contains(needle)) + { + return (true, Some("connection name fallback".to_string())); + } + + (false, None) +} + +fn first_executable_token(sql: &str) -> Option<&str> { + let bytes = sql.as_bytes(); + let mut i = 0; + + while i < bytes.len() { + while i < bytes.len() && bytes[i].is_ascii_whitespace() { + i += 1; + } + + if i + 1 < bytes.len() && bytes[i] == b'-' && bytes[i + 1] == b'-' { + i += 2; + while i < bytes.len() && bytes[i] != b'\n' { + i += 1; + } + continue; + } + + if i + 1 < bytes.len() && bytes[i] == b'/' && bytes[i + 1] == b'*' { + i += 2; + while i + 1 < bytes.len() && !(bytes[i] == b'*' && bytes[i + 1] == b'/') { + i += 1; + } + i = (i + 2).min(bytes.len()); + continue; + } + + break; + } + + let start = i; + while i < bytes.len() && (bytes[i].is_ascii_alphabetic() || bytes[i] == b'_') { + i += 1; + } + + (i > start).then_some(&sql[start..i]) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn comments_do_not_hide_read_token() { + assert_eq!(classify_sql("-- comment\nSELECT 1"), OperationClass::Read); + assert_eq!(classify_sql("/* DROP TABLE x */ SELECT 1"), OperationClass::Read); + } + + #[test] + fn classifies_write_and_ddl() { + assert_eq!( + classify_sql("update users set name = 'a'"), + OperationClass::Write + ); + assert_eq!(classify_sql("DROP TABLE users"), OperationClass::Ddl); + } + + #[test] + fn red_color_marks_production() { + let risk = risk_for("SELECT * FROM orders", "prod-main", Some("#ef4444")); + assert!(risk.is_production); + assert_eq!(risk.risk_level, RiskLevel::Low); + } +} From 99e5e6317d161f3c71404ee3dd6145c4d855563e Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 17:07:35 +0800 Subject: [PATCH 09/40] =?UTF-8?q?fix(sql-safety):=20=E4=BF=AE=E6=AD=A3?= =?UTF-8?q?=E5=8D=B1=E9=99=A9=20SQL=20=E5=88=86=E7=B1=BB=E4=B8=8E=E9=A3=8E?= =?UTF-8?q?=E9=99=A9=E4=B8=8A=E4=B8=8B=E6=96=87?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 扫描可执行 SQL token，避免 WITH、EXPLAIN ANALYZE 和多语句中的写入/DDL 被判为 Read - 引入 RiskContext 支持 environment_label，并保留 risk_for_connection 便捷调用 - 补充 SQL 安全分类与生产环境风险识别测试 BREAKING CHANGE: risk_for 现在接收 RiskContext；三参数便捷调用迁移为 risk_for_connection --- crates/dbx-core/src/sql_safety.rs | 170 ++++++++++++++++++++++++++---- 1 file changed, 152 insertions(+), 18 deletions(-) diff --git a/crates/dbx-core/src/sql_safety.rs b/crates/dbx-core/src/sql_safety.rs index 5f24ada2d..aedf92e3a 100644 --- a/crates/dbx-core/src/sql_safety.rs +++ b/crates/dbx-core/src/sql_safety.rs @@ -28,19 +28,47 @@ pub struct RiskMetadata { pub first_token: Option, } +#[derive(Debug, Clone, Copy)] +pub struct RiskContext<'a> { + pub connection_name: &'a str, + pub color: Option<&'a str>, + pub environment_label: Option<&'a str>, +} + +impl<'a> RiskContext<'a> { + pub fn new(connection_name: &'a str) -> Self { + Self { connection_name, color: None, environment_label: None } + } + + pub fn with_color(mut self, color: Option<&'a str>) -> Self { + self.color = color; + self + } + + pub fn with_environment_label(mut self, environment_label: Option<&'a str>) -> Self { + self.environment_label = environment_label; + self + } +} + pub fn classify_sql(sql: &str) -> OperationClass { - let token = first_executable_token(sql).map(|s| s.to_ascii_uppercase()); - match token.as_deref() { + let tokens = executable_tokens(sql); + if tokens.iter().any(|token| is_ddl_token(token)) { + return OperationClass::Ddl; + } + if tokens.iter().any(|token| is_write_token(token)) { + return OperationClass::Write; + } + + match tokens.first().map(String::as_str) { Some("SELECT" | "SHOW" | "DESCRIBE" | "EXPLAIN" | "WITH") => OperationClass::Read, - Some("INSERT" | "UPDATE" | "DELETE" | "MERGE" | "REPLACE") => OperationClass::Write, - Some("CREATE" | "ALTER" | "DROP" | "TRUNCATE" | "RENAME") => OperationClass::Ddl, _ => OperationClass::Unknown, } } -pub fn risk_for(sql: &str, connection_name: &str, color: Option<&str>) -> RiskMetadata { +pub fn risk_for(sql: &str, context: RiskContext<'_>) -> RiskMetadata { let operation_class = classify_sql(sql); - let (is_production, production_reason) = production_signal(connection_name, color); + let (is_production, production_reason) = production_signal(context); let risk_level = match (operation_class, is_production) { (OperationClass::Read, _) => RiskLevel::Low, (OperationClass::Write, false) => RiskLevel::Medium, @@ -58,22 +86,100 @@ pub fn risk_for(sql: &str, connection_name: &str, color: Option<&str>) -> RiskMe } } -fn production_signal(connection_name: &str, color: Option<&str>) -> (bool, Option) { - if matches!(color, Some("#ef4444")) { +pub fn risk_for_connection(sql: &str, connection_name: &str, color: Option<&str>) -> RiskMetadata { + risk_for(sql, RiskContext::new(connection_name).with_color(color)) +} + +fn production_signal(context: RiskContext<'_>) -> (bool, Option) { + if matches!(context.color, Some("#ef4444")) { return (true, Some("red connection color".to_string())); } - let name = connection_name.to_ascii_lowercase(); - if ["prod", "production", "live"] - .iter() - .any(|needle| name.contains(needle)) - { + if context.environment_label.is_some_and(contains_production_signal) { + return (true, Some("environment label".to_string())); + } + + if contains_production_signal(context.connection_name) { return (true, Some("connection name fallback".to_string())); } (false, None) } +fn contains_production_signal(value: &str) -> bool { + let value = value.to_ascii_lowercase(); + ["prod", "production", "live"].iter().any(|needle| value.contains(needle)) +} + +fn is_write_token(token: &str) -> bool { + matches!(token, "INSERT" | "UPDATE" | "DELETE" | "MERGE" | "REPLACE") +} + +fn is_ddl_token(token: &str) -> bool { + matches!(token, "CREATE" | "ALTER" | "DROP" | "TRUNCATE" | "RENAME") +} + +fn executable_tokens(sql: &str) -> Vec { + let bytes = sql.as_bytes(); + let mut i = 0; + let mut tokens = Vec::new(); + + while i < bytes.len() { + if bytes[i].is_ascii_whitespace() || bytes[i] == b';' { + i += 1; + continue; + } + + if i + 1 < bytes.len() && bytes[i] == b'-' && bytes[i + 1] == b'-' { + i += 2; + while i < bytes.len() && bytes[i] != b'\n' { + i += 1; + } + continue; + } + + if i + 1 < bytes.len() && bytes[i] == b'/' && bytes[i + 1] == b'*' { + i += 2; + while i + 1 < bytes.len() && !(bytes[i] == b'*' && bytes[i + 1] == b'/') { + i += 1; + } + i = (i + 2).min(bytes.len()); + continue; + } + + if matches!(bytes[i], b'\'' | b'"' | b'`') { + let quote = bytes[i]; + i += 1; + while i < bytes.len() { + if bytes[i] == quote { + if i + 1 < bytes.len() && bytes[i + 1] == quote { + i += 2; + continue; + } + i += 1; + break; + } + i += 1; + } + continue; + } + + if bytes[i].is_ascii_alphabetic() || bytes[i] == b'_' { + let start = i; + i += 1; + while i < bytes.len() && (bytes[i].is_ascii_alphanumeric() || bytes[i] == b'_') { + i += 1; + } + tokens.push(sql[start..i].to_ascii_uppercase()); + continue; + } + + i += 1; + } + + tokens +} + fn first_executable_token(sql: &str) -> Option<&str> { let bytes = sql.as_bytes(); let mut i = 0; @@ -123,17 +229,45 @@ mod tests { #[test] fn classifies_write_and_ddl() { - assert_eq!( - classify_sql("update users set name = 'a'"), - OperationClass::Write - ); + assert_eq!(classify_sql("update users set name = 'a'"), OperationClass::Write); assert_eq!(classify_sql("DROP TABLE users"), OperationClass::Ddl); } + #[test] + fn with_does_not_hide_write_or_ddl() { + assert_eq!( + classify_sql("WITH moved AS (DELETE FROM orders RETURNING *) SELECT * FROM moved"), + OperationClass::Write + ); + assert_eq!(classify_sql("WITH dropped AS (DROP TABLE old_orders) SELECT 1"), OperationClass::Ddl); + } + + #[test] + fn explain_analyze_write_is_write() { + assert_eq!(classify_sql("EXPLAIN ANALYZE UPDATE users SET name = 'a'"), OperationClass::Write); + } + + #[test] + fn dangerous_statement_in_multi_statement_sql_is_not_read() { + assert_eq!(classify_sql("SELECT * FROM users; DELETE FROM users WHERE id = 1"), OperationClass::Write); + assert_eq!(classify_sql("SHOW TABLES; DROP TABLE users"), OperationClass::Ddl); + } + #[test] fn red_color_marks_production() { - let risk = risk_for("SELECT * FROM orders", "prod-main", Some("#ef4444")); + let risk = risk_for_connection("SELECT * FROM orders", "prod-main", Some("#ef4444")); assert!(risk.is_production); assert_eq!(risk.risk_level, RiskLevel::Low); } + + #[test] + fn environment_label_marks_production() { + let risk = risk_for( + "UPDATE orders SET status = 'done'", + RiskContext { connection_name: "analytics", color: None, environment_label: Some("Production") }, + ); + assert!(risk.is_production); + assert_eq!(risk.production_reason.as_deref(), Some("environment label")); + assert_eq!(risk.risk_level, RiskLevel::High); + } } From 1909856467303cc9ef214d6e1f2bb46907a45894 Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 17:14:47 +0800 Subject: [PATCH 10/40] =?UTF-8?q?fix(sql-safety):=20=E4=BF=AE=E5=A4=8D=20S?= =?UTF-8?q?QL=20=E9=A3=8E=E9=99=A9=E5=A4=8D=E5=AE=A1=E9=97=AE=E9=A2=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 调整环境标签优先级，并阻止显式非生产标签回退 - 将无 WHERE/LIMIT 的 DELETE/UPDATE 标记为 critical - 跳过 PostgreSQL dollar quote，扫描 MySQL executable comment - 补充 sql_safety 覆盖测试 --- crates/dbx-core/src/sql_safety.rs | 207 ++++++++++++++++++++++++------ 1 file changed, 165 insertions(+), 42 deletions(-) diff --git a/crates/dbx-core/src/sql_safety.rs b/crates/dbx-core/src/sql_safety.rs index aedf92e3a..c8fb3d634 100644 --- a/crates/dbx-core/src/sql_safety.rs +++ b/crates/dbx-core/src/sql_safety.rs @@ -53,6 +53,10 @@ impl<'a> RiskContext<'a> { pub fn classify_sql(sql: &str) -> OperationClass { let tokens = executable_tokens(sql); + classify_tokens(&tokens) +} + +fn classify_tokens(tokens: &[String]) -> OperationClass { if tokens.iter().any(|token| is_ddl_token(token)) { return OperationClass::Ddl; } @@ -71,6 +75,7 @@ pub fn risk_for(sql: &str, context: RiskContext<'_>) -> RiskMetadata { let (is_production, production_reason) = production_signal(context); let risk_level = match (operation_class, is_production) { (OperationClass::Read, _) => RiskLevel::Low, + (OperationClass::Write, _) if has_unfiltered_destructive_write(sql) => RiskLevel::Critical, (OperationClass::Write, false) => RiskLevel::Medium, (OperationClass::Write, true) => RiskLevel::High, (OperationClass::Ddl, _) => RiskLevel::Critical, @@ -82,7 +87,7 @@ pub fn risk_for(sql: &str, context: RiskContext<'_>) -> RiskMetadata { risk_level, is_production, production_reason, - first_token: first_executable_token(sql).map(str::to_string), + first_token: first_executable_token(sql), } } @@ -91,12 +96,17 @@ pub fn risk_for_connection(sql: &str, connection_name: &str, color: Option<&str> } fn production_signal(context: RiskContext<'_>) -> (bool, Option) { - if matches!(context.color, Some("#ef4444")) { - return (true, Some("red connection color".to_string())); + if let Some(environment_label) = context.environment_label { + if contains_non_production_signal(environment_label) { + return (false, None); + } + if contains_production_signal(environment_label) { + return (true, Some("environment label".to_string())); + } } - if context.environment_label.is_some_and(contains_production_signal) { - return (true, Some("environment label".to_string())); + if matches!(context.color, Some("#ef4444")) { + return (true, Some("red connection color".to_string())); } if contains_production_signal(context.connection_name) { @@ -111,6 +121,27 @@ fn contains_production_signal(value: &str) -> bool { ["prod", "production", "live"].iter().any(|needle| value.contains(needle)) } +fn contains_non_production_signal(value: &str) -> bool { + let value = value.to_ascii_lowercase(); + [ + "dev", + "development", + "test", + "testing", + "qa", + "stage", + "staging", + "local", + "sandbox", + "non-prod", + "non-production", + "non production", + "nonprod", + ] + .iter() + .any(|needle| value.contains(needle)) +} + fn is_write_token(token: &str) -> bool { matches!(token, "INSERT" | "UPDATE" | "DELETE" | "MERGE" | "REPLACE") } @@ -119,13 +150,42 @@ fn is_ddl_token(token: &str) -> bool { matches!(token, "CREATE" | "ALTER" | "DROP" | "TRUNCATE" | "RENAME") } +fn has_unfiltered_destructive_write(sql: &str) -> bool { + executable_statements(sql).into_iter().any(|statement| { + statement.iter().enumerate().any(|(index, token)| { + if !matches!(token.as_str(), "DELETE" | "UPDATE") { + return false; + } + + !statement[index + 1..].iter().any(|token| matches!(token.as_str(), "WHERE" | "LIMIT")) + }) + }) +} + fn executable_tokens(sql: &str) -> Vec { + executable_statements(sql).into_iter().flatten().collect() +} + +fn executable_statements(sql: &str) -> Vec> { + let mut statements = Vec::new(); + let mut current = Vec::new(); + scan_executable_tokens(sql, &mut current, &mut statements); + push_statement(&mut current, &mut statements); + statements +} + +fn scan_executable_tokens(sql: &str, current: &mut Vec, statements: &mut Vec>) { let bytes = sql.as_bytes(); let mut i = 0; - let mut tokens = Vec::new(); while i < bytes.len() { - if bytes[i].is_ascii_whitespace() || bytes[i] == b';' { + if bytes[i].is_ascii_whitespace() { + i += 1; + continue; + } + + if bytes[i] == b';' { + push_statement(current, statements); i += 1; continue; } @@ -139,11 +199,29 @@ fn executable_tokens(sql: &str) -> Vec { } if i + 1 < bytes.len() && bytes[i] == b'/' && bytes[i + 1] == b'*' { - i += 2; - while i + 1 < bytes.len() && !(bytes[i] == b'*' && bytes[i + 1] == b'/') { - i += 1; + if i + 2 < bytes.len() && bytes[i + 2] == b'!' { + let content_start = i + 3; + let content_end = block_comment_end(bytes, content_start); + scan_executable_tokens(&sql[content_start..content_end], current, statements); + i = (content_end + 2).min(bytes.len()); + } else { + i += 2; + while i + 1 < bytes.len() && !(bytes[i] == b'*' && bytes[i + 1] == b'/') { + i += 1; + } + i = (i + 2).min(bytes.len()); + } + continue; + } + + if let Some(delimiter_len) = dollar_quote_delimiter_len(bytes, i) { + let delimiter = &sql[i..i + delimiter_len]; + i += delimiter_len; + if let Some(end) = sql[i..].find(delimiter) { + i += end + delimiter_len; + } else { + i = bytes.len(); } - i = (i + 2).min(bytes.len()); continue; } @@ -170,51 +248,54 @@ fn executable_tokens(sql: &str) -> Vec { while i < bytes.len() && (bytes[i].is_ascii_alphanumeric() || bytes[i] == b'_') { i += 1; } - tokens.push(sql[start..i].to_ascii_uppercase()); + current.push(sql[start..i].to_ascii_uppercase()); continue; } i += 1; } - - tokens } -fn first_executable_token(sql: &str) -> Option<&str> { - let bytes = sql.as_bytes(); - let mut i = 0; +fn push_statement(current: &mut Vec, statements: &mut Vec>) { + if !current.is_empty() { + statements.push(std::mem::take(current)); + } +} - while i < bytes.len() { - while i < bytes.len() && bytes[i].is_ascii_whitespace() { - i += 1; +fn block_comment_end(bytes: &[u8], mut i: usize) -> usize { + while i + 1 < bytes.len() { + if bytes[i] == b'*' && bytes[i + 1] == b'/' { + return i; } + i += 1; + } + bytes.len() +} - if i + 1 < bytes.len() && bytes[i] == b'-' && bytes[i + 1] == b'-' { - i += 2; - while i < bytes.len() && bytes[i] != b'\n' { - i += 1; - } - continue; - } - - if i + 1 < bytes.len() && bytes[i] == b'/' && bytes[i + 1] == b'*' { - i += 2; - while i + 1 < bytes.len() && !(bytes[i] == b'*' && bytes[i + 1] == b'/') { - i += 1; - } - i = (i + 2).min(bytes.len()); - continue; - } - - break; +fn dollar_quote_delimiter_len(bytes: &[u8], start: usize) -> Option { + if bytes.get(start) != Some(&b'$') { + return None; } - let start = i; - while i < bytes.len() && (bytes[i].is_ascii_alphabetic() || bytes[i] == b'_') { + let mut i = start + 1; + if bytes.get(i) == Some(&b'$') { + return Some(2); + } + + if !bytes.get(i).is_some_and(|byte| byte.is_ascii_alphabetic() || *byte == b'_') { + return None; + } + + i += 1; + while i < bytes.len() && (bytes[i].is_ascii_alphanumeric() || bytes[i] == b'_') { i += 1; } - (i > start).then_some(&sql[start..i]) + (bytes.get(i) == Some(&b'$')).then_some(i - start + 1) +} + +fn first_executable_token(sql: &str) -> Option { + executable_tokens(sql).into_iter().next() } #[cfg(test)] @@ -263,11 +344,53 @@ mod tests { #[test] fn environment_label_marks_production() { let risk = risk_for( - "UPDATE orders SET status = 'done'", + "UPDATE orders SET status = 'done' WHERE id = 1", RiskContext { connection_name: "analytics", color: None, environment_label: Some("Production") }, ); assert!(risk.is_production); assert_eq!(risk.production_reason.as_deref(), Some("environment label")); assert_eq!(risk.risk_level, RiskLevel::High); } + + #[test] + fn environment_label_overrides_color_and_name_fallback() { + let non_prod_label = risk_for( + "SELECT * FROM orders", + RiskContext { connection_name: "prod-main", color: Some("#ef4444"), environment_label: Some("Staging") }, + ); + assert!(!non_prod_label.is_production); + assert_eq!(non_prod_label.production_reason, None); + + let prod_label = risk_for( + "SELECT * FROM orders", + RiskContext { connection_name: "analytics", color: Some("#22c55e"), environment_label: Some("Production") }, + ); + assert!(prod_label.is_production); + assert_eq!(prod_label.production_reason.as_deref(), Some("environment label")); + } + + #[test] + fn destructive_writes_without_where_or_limit_are_critical() { + assert_eq!(risk_for("DELETE FROM users", RiskContext::new("dev")).risk_level, RiskLevel::Critical); + assert_eq!( + risk_for("UPDATE users SET active = false", RiskContext::new("dev")).risk_level, + RiskLevel::Critical + ); + assert_eq!(risk_for("DELETE FROM users WHERE id = 1", RiskContext::new("dev")).risk_level, RiskLevel::Medium); + assert_eq!(risk_for("TRUNCATE TABLE users", RiskContext::new("dev")).risk_level, RiskLevel::Critical); + } + + #[test] + fn postgresql_dollar_quotes_do_not_contribute_tokens() { + assert_eq!(classify_sql("SELECT $$ DELETE FROM users $$"), OperationClass::Read); + assert_eq!(classify_sql("SELECT $tag$ DROP TABLE users $tag$"), OperationClass::Read); + } + + #[test] + fn mysql_executable_comment_contributes_tokens() { + assert_eq!(classify_sql("/*!50000 DELETE FROM users */ SELECT 1"), OperationClass::Write); + let risk = risk_for("/*! UPDATE users SET active = false */", RiskContext::new("dev")); + assert_eq!(risk.operation_class, OperationClass::Write); + assert_eq!(risk.first_token.as_deref(), Some("UPDATE")); + } } From 99e79976209b6aee3434746dcb5ed8a68a9db594 Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 17:19:15 +0800 Subject: [PATCH 11/40] =?UTF-8?q?fix(sql=5Fsafety):=20=E4=BF=AE=E6=AD=A3?= =?UTF-8?q?=E5=AD=90=E6=9F=A5=E8=AF=A2=E8=BE=B9=E7=95=8C=E8=AF=AF=E5=88=A4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- crates/dbx-core/src/sql_safety.rs | 81 +++++++++++++++++++++++++++---- 1 file changed, 72 insertions(+), 9 deletions(-) diff --git a/crates/dbx-core/src/sql_safety.rs b/crates/dbx-core/src/sql_safety.rs index c8fb3d634..bc56db6d4 100644 --- a/crates/dbx-core/src/sql_safety.rs +++ b/crates/dbx-core/src/sql_safety.rs @@ -150,14 +150,22 @@ fn is_ddl_token(token: &str) -> bool { matches!(token, "CREATE" | "ALTER" | "DROP" | "TRUNCATE" | "RENAME") } +#[derive(Debug, Clone, PartialEq, Eq)] +struct SqlToken { + text: String, + depth: usize, +} + fn has_unfiltered_destructive_write(sql: &str) -> bool { - executable_statements(sql).into_iter().any(|statement| { + scanned_executable_statements(sql).into_iter().any(|statement| { statement.iter().enumerate().any(|(index, token)| { - if !matches!(token.as_str(), "DELETE" | "UPDATE") { + if !matches!(token.text.as_str(), "DELETE" | "UPDATE") { return false; } - !statement[index + 1..].iter().any(|token| matches!(token.as_str(), "WHERE" | "LIMIT")) + !statement[index + 1..] + .iter() + .any(|boundary| boundary.depth == token.depth && matches!(boundary.text.as_str(), "WHERE" | "LIMIT")) }) }) } @@ -167,14 +175,27 @@ fn executable_tokens(sql: &str) -> Vec { } fn executable_statements(sql: &str) -> Vec> { + scanned_executable_statements(sql) + .into_iter() + .map(|statement| statement.into_iter().map(|token| token.text).collect()) + .collect() +} + +fn scanned_executable_statements(sql: &str) -> Vec> { let mut statements = Vec::new(); let mut current = Vec::new(); - scan_executable_tokens(sql, &mut current, &mut statements); + let mut depth = 0; + scan_executable_tokens(sql, &mut current, &mut statements, &mut depth); push_statement(&mut current, &mut statements); statements } -fn scan_executable_tokens(sql: &str, current: &mut Vec, statements: &mut Vec>) { +fn scan_executable_tokens( + sql: &str, + current: &mut Vec, + statements: &mut Vec>, + depth: &mut usize, +) { let bytes = sql.as_bytes(); let mut i = 0; @@ -185,7 +206,21 @@ fn scan_executable_tokens(sql: &str, current: &mut Vec, statements: &mut } if bytes[i] == b';' { - push_statement(current, statements); + if *depth == 0 { + push_statement(current, statements); + } + i += 1; + continue; + } + + if bytes[i] == b'(' { + *depth += 1; + i += 1; + continue; + } + + if bytes[i] == b')' { + *depth = depth.saturating_sub(1); i += 1; continue; } @@ -202,7 +237,7 @@ fn scan_executable_tokens(sql: &str, current: &mut Vec, statements: &mut if i + 2 < bytes.len() && bytes[i + 2] == b'!' { let content_start = i + 3; let content_end = block_comment_end(bytes, content_start); - scan_executable_tokens(&sql[content_start..content_end], current, statements); + scan_executable_tokens(&sql[content_start..content_end], current, statements, depth); i = (content_end + 2).min(bytes.len()); } else { i += 2; @@ -248,7 +283,7 @@ fn scan_executable_tokens(sql: &str, current: &mut Vec, statements: &mut while i < bytes.len() && (bytes[i].is_ascii_alphanumeric() || bytes[i] == b'_') { i += 1; } - current.push(sql[start..i].to_ascii_uppercase()); + current.push(SqlToken { text: sql[start..i].to_ascii_uppercase(), depth: *depth }); continue; } @@ -256,7 +291,7 @@ fn scan_executable_tokens(sql: &str, current: &mut Vec, statements: &mut } } -fn push_statement(current: &mut Vec, statements: &mut Vec>) { +fn push_statement(current: &mut Vec, statements: &mut Vec>) { if !current.is_empty() { statements.push(std::mem::take(current)); } @@ -380,6 +415,34 @@ mod tests { assert_eq!(risk_for("TRUNCATE TABLE users", RiskContext::new("dev")).risk_level, RiskLevel::Critical); } + #[test] + fn destructive_writes_only_count_top_level_where_or_limit_as_boundaries() { + assert_eq!( + risk_for( + "DELETE FROM users USING (SELECT id FROM archived WHERE stale = true) old", + RiskContext::new("dev") + ) + .risk_level, + RiskLevel::Critical + ); + assert_eq!( + risk_for( + "UPDATE users SET active = false FROM (SELECT id FROM flags LIMIT 10) flags", + RiskContext::new("dev") + ) + .risk_level, + RiskLevel::Critical + ); + assert_eq!( + risk_for( + "DELETE FROM users WHERE id IN (SELECT user_id FROM archived WHERE stale = true)", + RiskContext::new("dev") + ) + .risk_level, + RiskLevel::Medium + ); + } + #[test] fn postgresql_dollar_quotes_do_not_contribute_tokens() { assert_eq!(classify_sql("SELECT $$ DELETE FROM users $$"), OperationClass::Read); From 81a6dce5acf1022e893642d278981b48ca970b54 Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 17:30:34 +0800 Subject: [PATCH 12/40] =?UTF-8?q?fix(sql=5Fsafety):=20=E4=BF=AE=E5=A4=8D?= =?UTF-8?q?=20CTE=20=E5=85=84=E5=BC=9F=E7=89=87=E6=AE=B5=E8=BE=B9=E7=95=8C?= =?UTF-8?q?=E8=AF=AF=E5=88=A4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - CTE 内 DELETE/UPDATE 离开当前片段后停止查找 WHERE/LIMIT - 补充 sibling CTE WHERE/LIMIT 不应降低风险等级的回归测试 --- crates/dbx-core/src/sql_safety.rs | 37 ++++++++++++++++++++++++++++--- 1 file changed, 34 insertions(+), 3 deletions(-) diff --git a/crates/dbx-core/src/sql_safety.rs b/crates/dbx-core/src/sql_safety.rs index bc56db6d4..09cf4f68a 100644 --- a/crates/dbx-core/src/sql_safety.rs +++ b/crates/dbx-core/src/sql_safety.rs @@ -163,13 +163,24 @@ fn has_unfiltered_destructive_write(sql: &str) -> bool { return false; } - !statement[index + 1..] - .iter() - .any(|boundary| boundary.depth == token.depth && matches!(boundary.text.as_str(), "WHERE" | "LIMIT")) + !has_same_fragment_boundary(&statement, index, token.depth) }) }) } +fn has_same_fragment_boundary(statement: &[SqlToken], destructive_write_index: usize, depth: usize) -> bool { + for boundary in &statement[destructive_write_index + 1..] { + if boundary.depth < depth { + break; + } + if boundary.depth == depth && matches!(boundary.text.as_str(), "WHERE" | "LIMIT") { + return true; + } + } + + false +} + fn executable_tokens(sql: &str) -> Vec { executable_statements(sql).into_iter().flatten().collect() } @@ -443,6 +454,26 @@ mod tests { ); } + #[test] + fn cte_destructive_writes_do_not_use_sibling_cte_boundaries() { + assert_eq!( + risk_for( + "WITH deleted AS (DELETE FROM users RETURNING id), scoped AS (SELECT id FROM audit WHERE id = 1) SELECT * FROM scoped", + RiskContext::new("dev") + ) + .risk_level, + RiskLevel::Critical + ); + assert_eq!( + risk_for( + "WITH updated AS (UPDATE users SET active = false RETURNING id), scoped AS (SELECT id FROM audit LIMIT 1) SELECT * FROM scoped", + RiskContext::new("dev") + ) + .risk_level, + RiskLevel::Critical + ); + } + #[test] fn postgresql_dollar_quotes_do_not_contribute_tokens() { assert_eq!(classify_sql("SELECT $$ DELETE FROM users $$"), OperationClass::Read); From 72690f6ebfccf7948fa43fc08a25b635f7fab0f0 Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 17:36:48 +0800 Subject: [PATCH 13/40] =?UTF-8?q?feat(core):=20=E5=A2=9E=E5=8A=A0=20handof?= =?UTF-8?q?f=20=E9=98=9F=E5=88=97=E5=AD=98=E5=82=A8=E6=A8=A1=E5=9E=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 新增 handoff 状态与队列项 DTO - 为 Storage 增加 handoffs 表和 pending 读写方法 - 补充 handoff 队列保存、排序和状态过滤单测 --- crates/dbx-core/src/handoff.rs | 65 ++++++++++++++++++++ crates/dbx-core/src/lib.rs | 1 + crates/dbx-core/src/storage.rs | 109 +++++++++++++++++++++++++++++++++ 3 files changed, 175 insertions(+) create mode 100644 crates/dbx-core/src/handoff.rs diff --git a/crates/dbx-core/src/handoff.rs b/crates/dbx-core/src/handoff.rs new file mode 100644 index 000000000..4f4d980c7 --- /dev/null +++ b/crates/dbx-core/src/handoff.rs @@ -0,0 +1,65 @@ +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +use crate::sql_safety::{OperationClass, RiskLevel}; + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "lowercase")] +pub enum HandoffStatus { + Queued, + Shown, + Approved, + Rejected, + Executed, + Failed, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct HandoffItem { + pub id: String, + pub created_at: DateTime, + pub created_by: String, + pub connection_name: String, + pub database: Option, + pub title: String, + pub description: Option, + pub sql: String, + pub operation_class: OperationClass, + pub risk_level: RiskLevel, + pub is_production: bool, + pub status: HandoffStatus, + pub result_summary: Option, + pub error: Option, +} + +impl HandoffItem { + pub fn queued( + connection_name: String, + database: Option, + title: String, + description: Option, + sql: String, + operation_class: OperationClass, + risk_level: RiskLevel, + is_production: bool, + ) -> Self { + Self { + id: Uuid::new_v4().to_string(), + created_at: Utc::now(), + created_by: "dbx-cli".to_string(), + connection_name, + database, + title, + description, + sql, + operation_class, + risk_level, + is_production, + status: HandoffStatus::Queued, + result_summary: None, + error: None, + } + } +} diff --git a/crates/dbx-core/src/lib.rs b/crates/dbx-core/src/lib.rs index 0636c7e43..6f309d441 100644 --- a/crates/dbx-core/src/lib.rs +++ b/crates/dbx-core/src/lib.rs @@ -4,6 +4,7 @@ pub mod connection; pub mod connection_secrets; pub mod db; pub mod external; +pub mod handoff; pub mod history; pub mod models; pub mod mongo_ops; diff --git a/crates/dbx-core/src/storage.rs b/crates/dbx-core/src/storage.rs index 27d588737..43686c03a 100644 --- a/crates/dbx-core/src/storage.rs +++ b/crates/dbx-core/src/storage.rs @@ -4,6 +4,7 @@ use std::str::FromStr; use sqlx::sqlite::{SqliteConnectOptions, SqlitePool, SqlitePoolOptions}; use crate::ai::{AiChatMessage, AiConfig, AiConversation}; +use crate::handoff::HandoffItem; use crate::history::HistoryEntry; use crate::models::connection::ConnectionConfig; use crate::saved_sql::{SavedSqlFile, SavedSqlFolder, SavedSqlLibrary}; @@ -84,6 +85,12 @@ const SCHEMA_STATEMENTS: &[&str] = &[ created_at TEXT NOT NULL DEFAULT '', updated_at TEXT NOT NULL DEFAULT '' )", + "CREATE TABLE IF NOT EXISTS handoffs ( + id TEXT PRIMARY KEY, + payload_json TEXT NOT NULL, + status TEXT NOT NULL, + created_at TEXT NOT NULL + )", ]; // --------------------------------------------------------------------------- @@ -242,6 +249,47 @@ impl Storage { } } +// --------------------------------------------------------------------------- +// Handoffs +// --------------------------------------------------------------------------- + +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())?; + + sqlx::query( + "INSERT OR REPLACE INTO handoffs (id, payload_json, status, created_at) \ + VALUES (?, ?, ?, ?)", + ) + .bind(&item.id) + .bind(json) + .bind(status) + .bind(item.created_at.to_rfc3339()) + .execute(&self.db) + .await + .map_err(|e| e.to_string())?; + + Ok(()) + } + + pub async fn load_pending_handoffs(&self) -> Result, String> { + let rows: Vec<(String,)> = sqlx::query_as( + "SELECT payload_json FROM handoffs \ + WHERE status IN ('queued', 'shown') \ + ORDER BY created_at ASC", + ) + .fetch_all(&self.db) + .await + .map_err(|e| e.to_string())?; + + rows.into_iter().map(|(json,)| serde_json::from_str(&json).map_err(|e| e.to_string())).collect() + } +} + // --------------------------------------------------------------------------- // AI Config // --------------------------------------------------------------------------- @@ -898,6 +946,67 @@ impl Storage { } } +#[cfg(test)] +mod handoff_tests { + use super::*; + use crate::handoff::{HandoffItem, HandoffStatus}; + use crate::sql_safety::{OperationClass, RiskLevel}; + + async fn open_temp_storage() -> Storage { + let path = std::env::temp_dir().join(format!("dbx-handoff-test-{}.db", uuid::Uuid::new_v4())); + Storage::open(&path).await.unwrap() + } + + fn queued_handoff(title: &str) -> HandoffItem { + HandoffItem::queued( + "prod-main".to_string(), + Some("app".to_string()), + title.to_string(), + Some("review write".to_string()), + "UPDATE users SET active = 0".to_string(), + OperationClass::Write, + RiskLevel::High, + true, + ) + } + + #[tokio::test] + async fn save_handoff_loads_pending_records_in_created_order() { + let storage = open_temp_storage().await; + let first = queued_handoff("first"); + let mut second = queued_handoff("second"); + second.created_at = first.created_at + chrono::Duration::seconds(1); + second.status = HandoffStatus::Shown; + + storage.save_handoff(&second).await.unwrap(); + storage.save_handoff(&first).await.unwrap(); + + let loaded = storage.load_pending_handoffs().await.unwrap(); + + assert_eq!(loaded.iter().map(|item| item.title.as_str()).collect::>(), vec!["first", "second"]); + assert_eq!(loaded[0].status, HandoffStatus::Queued); + assert_eq!(loaded[1].status, HandoffStatus::Shown); + assert_eq!(loaded[0].operation_class, OperationClass::Write); + assert!(loaded[0].is_production); + } + + #[tokio::test] + async fn load_pending_handoffs_excludes_terminal_statuses() { + let storage = open_temp_storage().await; + let queued = queued_handoff("queued"); + let mut executed = queued_handoff("executed"); + executed.status = HandoffStatus::Executed; + + storage.save_handoff(&queued).await.unwrap(); + storage.save_handoff(&executed).await.unwrap(); + + let loaded = storage.load_pending_handoffs().await.unwrap(); + + assert_eq!(loaded.len(), 1); + assert_eq!(loaded[0].id, queued.id); + } +} + // --------------------------------------------------------------------------- // Helpers // --------------------------------------------------------------------------- From 037e082104296324b2be93991819e8751d88f979 Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 17:42:42 +0800 Subject: [PATCH 14/40] =?UTF-8?q?feat(handoff):=20=E5=A2=9E=E5=8A=A0?= =?UTF-8?q?=E8=BF=9E=E6=8E=A5=20ID=20=E5=92=8C=E7=A8=B3=E5=AE=9A=E6=8E=92?= =?UTF-8?q?=E5=BA=8F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 为 HandoffItem 增加 connection_id，并保留 connection_name 用于展示 - 为 handoffs 增加自增 seq 并迁移旧表结构 - load_pending_handoffs 按 created_at 与 seq 稳定 FIFO 排序 - 补充同 created_at FIFO、seq 结构与 connectionId 序列化测试 --- crates/dbx-core/src/handoff.rs | 4 ++ crates/dbx-core/src/storage.rs | 93 ++++++++++++++++++++++++++++++++-- 2 files changed, 93 insertions(+), 4 deletions(-) diff --git a/crates/dbx-core/src/handoff.rs b/crates/dbx-core/src/handoff.rs index 4f4d980c7..b239647f2 100644 --- a/crates/dbx-core/src/handoff.rs +++ b/crates/dbx-core/src/handoff.rs @@ -21,6 +21,8 @@ pub struct HandoffItem { pub id: String, pub created_at: DateTime, pub created_by: String, + #[serde(default)] + pub connection_id: String, pub connection_name: String, pub database: Option, pub title: String, @@ -36,6 +38,7 @@ pub struct HandoffItem { impl HandoffItem { pub fn queued( + connection_id: String, connection_name: String, database: Option, title: String, @@ -49,6 +52,7 @@ impl HandoffItem { id: Uuid::new_v4().to_string(), created_at: Utc::now(), created_by: "dbx-cli".to_string(), + connection_id, connection_name, database, title, diff --git a/crates/dbx-core/src/storage.rs b/crates/dbx-core/src/storage.rs index 43686c03a..444792949 100644 --- a/crates/dbx-core/src/storage.rs +++ b/crates/dbx-core/src/storage.rs @@ -86,7 +86,8 @@ const SCHEMA_STATEMENTS: &[&str] = &[ updated_at TEXT NOT NULL DEFAULT '' )", "CREATE TABLE IF NOT EXISTS handoffs ( - id TEXT PRIMARY KEY, + seq INTEGER PRIMARY KEY AUTOINCREMENT, + id TEXT NOT NULL UNIQUE, payload_json TEXT NOT NULL, status TEXT NOT NULL, created_at TEXT NOT NULL @@ -108,6 +109,7 @@ impl Storage { sqlx::query(statement).execute(&pool).await.map_err(|e| e.to_string())?; } ensure_history_columns(&pool).await?; + ensure_handoffs_sequence(&pool).await?; Ok(Self { db: pool }) } @@ -141,6 +143,48 @@ async fn ensure_history_columns(pool: &SqlitePool) -> Result<(), String> { Ok(()) } +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())?; + + if seq_columns > 0 { + return Ok(()); + } + + let mut tx = pool.begin().await.map_err(|e| e.to_string())?; + sqlx::query( + "CREATE TABLE handoffs_migration ( + seq INTEGER PRIMARY KEY AUTOINCREMENT, + id TEXT NOT NULL UNIQUE, + payload_json TEXT NOT NULL, + status TEXT NOT NULL, + created_at TEXT NOT NULL + )", + ) + .execute(&mut *tx) + .await + .map_err(|e| e.to_string())?; + sqlx::query( + "INSERT INTO handoffs_migration (id, payload_json, status, created_at) + SELECT id, payload_json, status, created_at + FROM handoffs + ORDER BY created_at ASC, rowid ASC", + ) + .execute(&mut *tx) + .await + .map_err(|e| e.to_string())?; + sqlx::query("DROP TABLE handoffs").execute(&mut *tx).await.map_err(|e| e.to_string())?; + sqlx::query("ALTER TABLE handoffs_migration RENAME TO handoffs") + .execute(&mut *tx) + .await + .map_err(|e| e.to_string())?; + tx.commit().await.map_err(|e| e.to_string())?; + Ok(()) +} + // --------------------------------------------------------------------------- // History // --------------------------------------------------------------------------- @@ -262,8 +306,12 @@ impl Storage { .ok_or_else(|| "Failed to serialize handoff status".to_string())?; sqlx::query( - "INSERT OR REPLACE INTO handoffs (id, payload_json, status, created_at) \ - VALUES (?, ?, ?, ?)", + "INSERT INTO handoffs (id, payload_json, status, created_at) \ + VALUES (?, ?, ?, ?) \ + ON CONFLICT(id) DO UPDATE SET \ + payload_json = excluded.payload_json, \ + status = excluded.status, \ + created_at = excluded.created_at", ) .bind(&item.id) .bind(json) @@ -280,7 +328,7 @@ impl Storage { let rows: Vec<(String,)> = sqlx::query_as( "SELECT payload_json FROM handoffs \ WHERE status IN ('queued', 'shown') \ - ORDER BY created_at ASC", + ORDER BY created_at ASC, seq ASC", ) .fetch_all(&self.db) .await @@ -959,6 +1007,7 @@ mod handoff_tests { fn queued_handoff(title: &str) -> HandoffItem { HandoffItem::queued( + "prod-main-id".to_string(), "prod-main".to_string(), Some("app".to_string()), title.to_string(), @@ -990,6 +1039,42 @@ mod handoff_tests { assert!(loaded[0].is_production); } + #[tokio::test] + async fn load_pending_handoffs_keeps_fifo_order_for_matching_created_at() { + let storage = open_temp_storage().await; + let first = queued_handoff("first"); + let mut second = queued_handoff("second"); + second.created_at = first.created_at; + + storage.save_handoff(&first).await.unwrap(); + storage.save_handoff(&second).await.unwrap(); + + let loaded = storage.load_pending_handoffs().await.unwrap(); + + assert_eq!(loaded.iter().map(|item| item.title.as_str()).collect::>(), vec!["first", "second"]); + } + + #[tokio::test] + 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(); + + assert!(columns.iter().any(|(name,)| name == "seq")); + } + + #[tokio::test] + async fn handoff_item_serializes_connection_id_and_display_name() { + let item = queued_handoff("serialize"); + + let value = serde_json::to_value(&item).unwrap(); + + assert_eq!(item.connection_id, "prod-main-id"); + assert_eq!(value["connectionId"], "prod-main-id"); + assert_eq!(value["connectionName"], "prod-main"); + } + #[tokio::test] async fn load_pending_handoffs_excludes_terminal_statuses() { let storage = open_temp_storage().await; From 78be118c58f7db36d14cdfa82c734fac888fbac8 Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 17:52:56 +0800 Subject: [PATCH 15/40] =?UTF-8?q?feat(schema):=20=E6=B7=BB=E5=8A=A0=20sche?= =?UTF-8?q?ma=20snapshot=20=E7=BC=96=E6=8E=92?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 新增 schema snapshot DTO，标准化连接、数据库、driver 和采集时间信息 - 复用现有 schema core 函数采集表、视图、列、索引、外键和触发器 - 添加 SQLite 集成测试覆盖表/视图与元数据快照输出 --- crates/dbx-core/src/lib.rs | 1 + crates/dbx-core/src/schema_snapshot.rs | 119 +++++++++++++++++++++++ crates/dbx-core/tests/schema_snapshot.rs | 99 +++++++++++++++++++ 3 files changed, 219 insertions(+) create mode 100644 crates/dbx-core/src/schema_snapshot.rs create mode 100644 crates/dbx-core/tests/schema_snapshot.rs diff --git a/crates/dbx-core/src/lib.rs b/crates/dbx-core/src/lib.rs index 6f309d441..78dd76b02 100644 --- a/crates/dbx-core/src/lib.rs +++ b/crates/dbx-core/src/lib.rs @@ -14,6 +14,7 @@ pub mod query_cancel; pub mod redis_ops; pub mod saved_sql; pub mod schema; +pub mod schema_snapshot; pub mod sql; pub mod sql_safety; pub mod storage; diff --git a/crates/dbx-core/src/schema_snapshot.rs b/crates/dbx-core/src/schema_snapshot.rs new file mode 100644 index 000000000..3ad9eff92 --- /dev/null +++ b/crates/dbx-core/src/schema_snapshot.rs @@ -0,0 +1,119 @@ +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; + +use crate::connection::AppState; +use crate::models::connection::{ConnectionConfig, DatabaseType}; +use crate::{schema, types}; + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct TableSnapshot { + pub name: String, + pub table_type: String, + pub comment: Option, + pub columns: Vec, + pub indexes: Vec, + pub foreign_keys: Vec, + pub triggers: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct SchemaSnapshot { + pub connection_id: String, + pub connection_name: String, + pub database: Option, + pub database_type: DatabaseType, + pub driver_profile: Option, + pub captured_at: DateTime, + pub databases: Vec, + pub schemas: Vec, + pub tables: Vec, + pub views: Vec, +} + +pub async fn snapshot( + state: &AppState, + connection_id: &str, + database: Option<&str>, + schema_name: Option<&str>, +) -> Result { + let config = { + let configs = state.configs.lock().await; + configs.get(connection_id).cloned().ok_or("Connection config not found")? + }; + + let db = snapshot_database(&config, database); + let databases = schema::list_databases_core(state, connection_id).await.unwrap_or_default(); + let schemas = if db.is_empty() { + Vec::new() + } else { + schema::list_schemas_core(state, connection_id, &db).await.unwrap_or_default() + }; + let effective_schema = schema_name.or_else(|| schemas.first().map(String::as_str)).unwrap_or(""); + let table_infos = if db.is_empty() { + Vec::new() + } else { + schema::list_tables_core(state, connection_id, &db, effective_schema).await.unwrap_or_default() + }; + + let mut tables = Vec::new(); + let mut views = Vec::new(); + for table in table_infos { + let table_snapshot = TableSnapshot { + columns: schema::get_columns_core(state, connection_id, &db, effective_schema, &table.name) + .await + .unwrap_or_default(), + indexes: schema::list_indexes_core(state, connection_id, &db, effective_schema, &table.name) + .await + .unwrap_or_default(), + foreign_keys: schema::list_foreign_keys_core(state, connection_id, &db, effective_schema, &table.name) + .await + .unwrap_or_default(), + triggers: schema::list_triggers_core(state, connection_id, &db, effective_schema, &table.name) + .await + .unwrap_or_default(), + name: table.name, + table_type: table.table_type, + comment: table.comment, + }; + + if is_view(&table_snapshot.table_type) { + views.push(table_snapshot); + } else { + tables.push(table_snapshot); + } + } + + Ok(SchemaSnapshot { + connection_id: config.id, + connection_name: config.name, + database: (!db.is_empty()).then_some(db), + database_type: config.db_type, + driver_profile: config.driver_profile, + captured_at: Utc::now(), + databases, + schemas, + tables, + views, + }) +} + +fn snapshot_database(config: &ConnectionConfig, requested_database: Option<&str>) -> String { + requested_database + .or_else(|| config.effective_database()) + .or_else(|| embedded_default_database(&config.db_type)) + .unwrap_or_default() + .to_string() +} + +fn embedded_default_database(db_type: &DatabaseType) -> Option<&'static str> { + match db_type { + DatabaseType::Sqlite | DatabaseType::DuckDb => Some("main"), + _ => None, + } +} + +fn is_view(table_type: &str) -> bool { + table_type.eq_ignore_ascii_case("view") +} diff --git a/crates/dbx-core/tests/schema_snapshot.rs b/crates/dbx-core/tests/schema_snapshot.rs new file mode 100644 index 000000000..d4f2ade3a --- /dev/null +++ b/crates/dbx-core/tests/schema_snapshot.rs @@ -0,0 +1,99 @@ +use std::str::FromStr; + +use dbx_core::connection::{AppState, PoolKind}; +use dbx_core::models::connection::{default_ssh_connect_timeout_secs, ConnectionConfig, DatabaseType}; +use dbx_core::schema_snapshot::snapshot; +use dbx_core::storage::Storage; +use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions}; + +fn sqlite_config(path: &std::path::Path) -> ConnectionConfig { + ConnectionConfig { + id: "sqlite-fixture".to_string(), + name: "SQLite Fixture".to_string(), + db_type: DatabaseType::Sqlite, + driver_profile: Some("builtin-sqlite".to_string()), + driver_label: None, + url_params: None, + host: path.display().to_string(), + port: 0, + username: String::new(), + password: String::new(), + database: None, + color: None, + ssh_enabled: false, + ssh_host: String::new(), + ssh_port: 22, + ssh_user: String::new(), + ssh_password: String::new(), + ssh_key_path: String::new(), + ssh_key_passphrase: String::new(), + ssh_expose_lan: false, + ssh_connect_timeout_secs: default_ssh_connect_timeout_secs(), + ssl: false, + sysdba: false, + connection_string: None, + jdbc_driver_class: None, + jdbc_driver_paths: Vec::new(), + } +} + +async fn create_sqlite_fixture(path: &std::path::Path) { + let url = format!("sqlite:{}?mode=rwc", path.display()); + let options = SqliteConnectOptions::from_str(&url).unwrap().create_if_missing(true); + let pool = SqlitePoolOptions::new().max_connections(1).connect_with(options).await.unwrap(); + + sqlx::query("PRAGMA foreign_keys = ON").execute(&pool).await.unwrap(); + sqlx::query("CREATE TABLE teams (id INTEGER PRIMARY KEY, name TEXT NOT NULL)").execute(&pool).await.unwrap(); + sqlx::query( + "CREATE TABLE users (id INTEGER PRIMARY KEY, team_id INTEGER NOT NULL, email TEXT NOT NULL UNIQUE, FOREIGN KEY(team_id) REFERENCES teams(id))", + ) + .execute(&pool) + .await + .unwrap(); + sqlx::query("CREATE INDEX idx_users_team_id ON users(team_id)").execute(&pool).await.unwrap(); + sqlx::query("CREATE TRIGGER trg_users_ai AFTER INSERT ON users BEGIN SELECT 1; END").execute(&pool).await.unwrap(); + sqlx::query("CREATE VIEW active_users AS SELECT id, email FROM users").execute(&pool).await.unwrap(); + + pool.close().await; +} + +async fn open_state() -> AppState { + let storage_path = std::env::temp_dir().join(format!("dbx-schema-snapshot-storage-{}.db", uuid::Uuid::new_v4())); + AppState::new(Storage::open(&storage_path).await.unwrap()) +} + +#[tokio::test] +async fn snapshot_standardizes_sqlite_tables_views_and_metadata() { + let data_path = std::env::temp_dir().join(format!("dbx-schema-snapshot-data-{}.db", uuid::Uuid::new_v4())); + create_sqlite_fixture(&data_path).await; + + let state = open_state().await; + let config = sqlite_config(&data_path); + let pool = dbx_core::db::sqlite::connect_path(&data_path.display().to_string()).await.unwrap(); + state.configs.lock().await.insert(config.id.clone(), config.clone()); + state.connections.lock().await.insert(config.id.clone(), PoolKind::Sqlite(pool)); + + let snapshot = snapshot(&state, &config.id, None, None).await.unwrap(); + + assert_eq!(snapshot.connection_id, "sqlite-fixture"); + assert_eq!(snapshot.connection_name, "SQLite Fixture"); + assert_eq!(snapshot.database.as_deref(), Some("main")); + assert_eq!(snapshot.database_type, DatabaseType::Sqlite); + assert_eq!(snapshot.driver_profile.as_deref(), Some("builtin-sqlite")); + let now = chrono::Utc::now(); + assert!(snapshot.captured_at <= now); + assert!(snapshot.captured_at > now - chrono::Duration::seconds(5)); + assert_eq!(snapshot.databases.iter().map(|db| db.name.as_str()).collect::>(), vec!["main"]); + + let table_names = snapshot.tables.iter().map(|table| table.name.as_str()).collect::>(); + assert_eq!(table_names, vec!["teams", "users"]); + let view_names = snapshot.views.iter().map(|view| view.name.as_str()).collect::>(); + assert_eq!(view_names, vec!["active_users"]); + + let users = snapshot.tables.iter().find(|table| table.name == "users").unwrap(); + assert_eq!(users.table_type, "BASE TABLE"); + assert!(users.columns.iter().any(|column| column.name == "email" && !column.is_nullable)); + assert!(users.indexes.iter().any(|index| index.name == "idx_users_team_id" && index.columns == vec!["team_id"])); + assert!(users.foreign_keys.iter().any(|fk| fk.column == "team_id" && fk.ref_table == "teams")); + assert!(users.triggers.iter().any(|trigger| trigger.name == "trg_users_ai" && trigger.event == "INSERT")); +} From 93f99039f1a587335e2e7f849d1e700fb9a88d79 Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 17:58:37 +0800 Subject: [PATCH 16/40] =?UTF-8?q?fix(schema=5Fsnapshot):=20=E4=BF=AE?= =?UTF-8?q?=E5=A4=8D=E5=BF=AB=E7=85=A7=E8=A1=A8=E8=A7=86=E5=9B=BE=E5=BD=92?= =?UTF-8?q?=E4=B8=80=E4=B8=8E=E9=94=99=E8=AF=AF=E4=BC=A0=E6=92=AD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 移除独立 views 字段，统一通过 tables 返回表和视图 - 为逐表元数据采集增加受限并发 helper - 传播 schema core 错误并补充上下文信息 - 对缺少必需 database 的连接返回明确错误 - 补充 SQLite 视图、错误传播和无 database 场景单测 --- crates/dbx-core/src/schema_snapshot.rs | 123 ++++++++++++++++------- crates/dbx-core/tests/schema_snapshot.rs | 40 +++++++- 2 files changed, 121 insertions(+), 42 deletions(-) diff --git a/crates/dbx-core/src/schema_snapshot.rs b/crates/dbx-core/src/schema_snapshot.rs index 3ad9eff92..1732e65a4 100644 --- a/crates/dbx-core/src/schema_snapshot.rs +++ b/crates/dbx-core/src/schema_snapshot.rs @@ -1,10 +1,13 @@ use chrono::{DateTime, Utc}; +use futures::{stream, StreamExt, TryStreamExt}; use serde::{Deserialize, Serialize}; use crate::connection::AppState; use crate::models::connection::{ConnectionConfig, DatabaseType}; use crate::{schema, types}; +const TABLE_METADATA_CONCURRENCY: usize = 4; + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct TableSnapshot { @@ -29,7 +32,6 @@ pub struct SchemaSnapshot { pub databases: Vec, pub schemas: Vec, pub tables: Vec, - pub views: Vec, } pub async fn snapshot( @@ -43,47 +45,27 @@ pub async fn snapshot( configs.get(connection_id).cloned().ok_or("Connection config not found")? }; - let db = snapshot_database(&config, database); - let databases = schema::list_databases_core(state, connection_id).await.unwrap_or_default(); + let db = snapshot_database(&config, database)?; + let databases = schema::list_databases_core(state, connection_id) + .await + .map_err(|err| format!("Failed to list databases: {err}"))?; let schemas = if db.is_empty() { Vec::new() } else { - schema::list_schemas_core(state, connection_id, &db).await.unwrap_or_default() + schema::list_schemas_core(state, connection_id, &db) + .await + .map_err(|err| format!("Failed to list schemas for database '{db}': {err}"))? }; let effective_schema = schema_name.or_else(|| schemas.first().map(String::as_str)).unwrap_or(""); let table_infos = if db.is_empty() { Vec::new() } else { - schema::list_tables_core(state, connection_id, &db, effective_schema).await.unwrap_or_default() + schema::list_tables_core(state, connection_id, &db, effective_schema) + .await + .map_err(|err| format!("Failed to list tables for database '{db}' schema '{effective_schema}': {err}"))? }; - let mut tables = Vec::new(); - let mut views = Vec::new(); - for table in table_infos { - let table_snapshot = TableSnapshot { - columns: schema::get_columns_core(state, connection_id, &db, effective_schema, &table.name) - .await - .unwrap_or_default(), - indexes: schema::list_indexes_core(state, connection_id, &db, effective_schema, &table.name) - .await - .unwrap_or_default(), - foreign_keys: schema::list_foreign_keys_core(state, connection_id, &db, effective_schema, &table.name) - .await - .unwrap_or_default(), - triggers: schema::list_triggers_core(state, connection_id, &db, effective_schema, &table.name) - .await - .unwrap_or_default(), - name: table.name, - table_type: table.table_type, - comment: table.comment, - }; - - if is_view(&table_snapshot.table_type) { - views.push(table_snapshot); - } else { - tables.push(table_snapshot); - } - } + let tables = collect_table_snapshots(state, connection_id, &db, effective_schema, table_infos).await?; Ok(SchemaSnapshot { connection_id: config.id, @@ -95,16 +77,71 @@ pub async fn snapshot( databases, schemas, tables, - views, }) } -fn snapshot_database(config: &ConnectionConfig, requested_database: Option<&str>) -> String { - requested_database +async fn collect_table_snapshots( + state: &AppState, + connection_id: &str, + database: &str, + schema_name: &str, + table_infos: Vec, +) -> Result, String> { + stream::iter(table_infos) + .map(|table| async move { table_snapshot(state, connection_id, database, schema_name, table).await }) + .buffered(TABLE_METADATA_CONCURRENCY) + .try_collect() + .await +} + +async fn table_snapshot( + state: &AppState, + connection_id: &str, + database: &str, + schema_name: &str, + table: types::TableInfo, +) -> Result { + let table_name = table.name.clone(); + let columns = schema::get_columns_core(state, connection_id, database, schema_name, &table_name) + .await + .map_err(|err| format!("Failed to list columns for table '{table_name}': {err}"))?; + let indexes = schema::list_indexes_core(state, connection_id, database, schema_name, &table_name) + .await + .map_err(|err| format!("Failed to list indexes for table '{table_name}': {err}"))?; + let foreign_keys = schema::list_foreign_keys_core(state, connection_id, database, schema_name, &table_name) + .await + .map_err(|err| format!("Failed to list foreign keys for table '{table_name}': {err}"))?; + let triggers = schema::list_triggers_core(state, connection_id, database, schema_name, &table_name) + .await + .map_err(|err| format!("Failed to list triggers for table '{table_name}': {err}"))?; + + Ok(TableSnapshot { + name: table.name, + table_type: table.table_type, + comment: table.comment, + columns, + indexes, + foreign_keys, + triggers, + }) +} + +fn snapshot_database(config: &ConnectionConfig, requested_database: Option<&str>) -> Result { + let database = requested_database + .map(str::trim) + .filter(|database| !database.is_empty()) .or_else(|| config.effective_database()) .or_else(|| embedded_default_database(&config.db_type)) - .unwrap_or_default() - .to_string() + .map(str::to_string); + + match database { + Some(database) => Ok(database), + None if requires_database(&config.db_type) => Err(format!( + "Database is required for schema snapshot for connection '{}' ({:?})", + config.id, config.db_type + )), + None => Ok(String::new()), + } } fn embedded_default_database(db_type: &DatabaseType) -> Option<&'static str> { @@ -114,6 +151,14 @@ fn embedded_default_database(db_type: &DatabaseType) -> Option<&'static str> { } } -fn is_view(table_type: &str) -> bool { - table_type.eq_ignore_ascii_case("view") +fn requires_database(db_type: &DatabaseType) -> bool { + matches!( + db_type, + DatabaseType::Mysql + | DatabaseType::Doris + | DatabaseType::StarRocks + | DatabaseType::ClickHouse + | DatabaseType::MongoDb + | DatabaseType::Jdbc + ) } diff --git a/crates/dbx-core/tests/schema_snapshot.rs b/crates/dbx-core/tests/schema_snapshot.rs index d4f2ade3a..54c7500ab 100644 --- a/crates/dbx-core/tests/schema_snapshot.rs +++ b/crates/dbx-core/tests/schema_snapshot.rs @@ -86,9 +86,8 @@ async fn snapshot_standardizes_sqlite_tables_views_and_metadata() { assert_eq!(snapshot.databases.iter().map(|db| db.name.as_str()).collect::>(), vec!["main"]); let table_names = snapshot.tables.iter().map(|table| table.name.as_str()).collect::>(); - assert_eq!(table_names, vec!["teams", "users"]); - let view_names = snapshot.views.iter().map(|view| view.name.as_str()).collect::>(); - assert_eq!(view_names, vec!["active_users"]); + assert_eq!(table_names, vec!["active_users", "teams", "users"]); + assert!(serde_json::to_value(&snapshot).unwrap().get("views").is_none()); let users = snapshot.tables.iter().find(|table| table.name == "users").unwrap(); assert_eq!(users.table_type, "BASE TABLE"); @@ -96,4 +95,39 @@ async fn snapshot_standardizes_sqlite_tables_views_and_metadata() { assert!(users.indexes.iter().any(|index| index.name == "idx_users_team_id" && index.columns == vec!["team_id"])); assert!(users.foreign_keys.iter().any(|fk| fk.column == "team_id" && fk.ref_table == "teams")); assert!(users.triggers.iter().any(|trigger| trigger.name == "trg_users_ai" && trigger.event == "INSERT")); + + let active_users = snapshot.tables.iter().find(|table| table.name == "active_users").unwrap(); + assert_eq!(active_users.table_type, "VIEW"); + assert!(active_users.columns.iter().any(|column| column.name == "email")); +} + +#[tokio::test] +async fn snapshot_propagates_schema_core_errors() { + let data_path = std::env::temp_dir().join(format!("dbx-schema-snapshot-missing-{}.db", uuid::Uuid::new_v4())); + let state = open_state().await; + let config = sqlite_config(&data_path); + state.configs.lock().await.insert(config.id.clone(), config.clone()); + + let err = snapshot(&state, &config.id, None, None).await.unwrap_err(); + + assert!(err.contains("Failed to list databases"), "{err}"); + assert!(err.contains("Connection not found"), "{err}"); +} + +#[tokio::test] +async fn snapshot_requires_database_for_database_scoped_connections_without_default() { + let state = open_state().await; + let mut config = sqlite_config(std::path::Path::new("unused")); + config.id = "mysql-no-db".to_string(); + config.name = "MySQL without DB".to_string(); + config.db_type = DatabaseType::Mysql; + config.driver_profile = None; + config.host = "127.0.0.1".to_string(); + config.port = 3306; + state.configs.lock().await.insert(config.id.clone(), config.clone()); + + let err = snapshot(&state, &config.id, None, None).await.unwrap_err(); + + assert!(err.contains("Database is required"), "{err}"); + assert!(err.contains("mysql-no-db"), "{err}"); } From f48868006cdee0dfce2f3338c9c261d54c2b3407 Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 18:29:19 +0800 Subject: [PATCH 17/40] =?UTF-8?q?feat(dbx-cli):=20=E5=A2=9E=E5=8A=A0=20run?= =?UTF-8?q?time=20discovery=20=E4=B8=8E=E5=91=BD=E4=BB=A4=E5=88=86?= =?UTF-8?q?=E5=8F=91?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 增加 dbx-cli 参数解析与 8 个命令入口 - 实现 runtime discovery client 与 headless fallback - 更新 CLI 依赖和锁文件 - 纳入 cli/storage 的格式化 diff --- Cargo.lock | 1 + crates/dbx-cli/Cargo.toml | 3 + crates/dbx-cli/src/commands.rs | 213 +++++++++++++++++++++++++-- crates/dbx-cli/src/runtime_client.rs | 38 ++++- crates/dbx-core/src/cli.rs | 8 +- crates/dbx-core/src/storage.rs | 15 +- 6 files changed, 253 insertions(+), 25 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 27dfd5fef..2b20b22c7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1682,6 +1682,7 @@ dependencies = [ "reqwest 0.12.28", "serde", "serde_json", + "tempfile", "tokio", "uuid", ] diff --git a/crates/dbx-cli/Cargo.toml b/crates/dbx-cli/Cargo.toml index d272aaa4b..a81c00834 100644 --- a/crates/dbx-cli/Cargo.toml +++ b/crates/dbx-cli/Cargo.toml @@ -15,3 +15,6 @@ serde_json = "1.0" tokio = { version = "1", features = ["full"] } reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] } uuid = { version = "1", features = ["v4", "serde"] } + +[dev-dependencies] +tempfile = "3" diff --git a/crates/dbx-cli/src/commands.rs b/crates/dbx-cli/src/commands.rs index 9ab5ab8d0..e97350dd6 100644 --- a/crates/dbx-cli/src/commands.rs +++ b/crates/dbx-cli/src/commands.rs @@ -1,26 +1,219 @@ use dbx_core::cli::{fail, ok, CliEnvelope, CliErrorCode, CliSource}; pub(crate) async fn run(args: Vec) -> Result<(), CliEnvelope<()>> { - let output = match args.as_slice() { - [cmd, ..] if cmd == "context" => context().await, - [cmd, ..] if cmd == "selection" => runtime_required("dbx selection requires DBX GUI runtime."), - [cmd, sub, ..] if cmd == "result" && sub == "current" => { - runtime_required("dbx result current requires DBX GUI runtime.") - } - _ => ok(CliSource::Headless, serde_json::json!({ "runtime": "headless" })), - }; - + let output = dispatch(args).await; println!("{}", serde_json::to_string_pretty(&output).unwrap()); + + if matches!(output, CliEnvelope::Failure { .. }) { + std::process::exit(1); + } + Ok(()) } -async fn context() -> CliEnvelope { +pub(crate) async fn dispatch(args: Vec) -> CliEnvelope { + let parsed = match parse_args(args) { + Ok(parsed) => parsed, + Err(err) => return err, + }; + + match parsed.positionals.as_slice() { + [cmd, rest @ ..] if cmd == "context" => context(rest).await, + [cmd, sub, rest @ ..] if cmd == "conn" && sub == "list" => conn_list(rest).await, + [cmd, sub, name, rest @ ..] if cmd == "conn" && sub == "show" => conn_show(name, rest).await, + [cmd, sub, rest @ ..] if cmd == "schema" && sub == "snapshot" => schema_snapshot(rest).await, + [cmd, rest @ ..] if cmd == "safe-query" => safe_query(rest).await, + [cmd, rest @ ..] if cmd == "handoff" => handoff(rest).await, + [cmd, rest @ ..] if cmd == "selection" => selection(rest).await, + [cmd, sub, rest @ ..] if cmd == "result" && sub == "current" => result_current(rest).await, + _ => fail(CliSource::Headless, CliErrorCode::InternalError, "Unknown command", false), + } +} + +struct ParsedArgs { + positionals: Vec, +} + +fn parse_args(args: Vec) -> Result> { + let mut positionals = Vec::new(); + let mut index = 0; + + while index < args.len() { + if args[index] == "--format" { + let Some(format) = args.get(index + 1) else { + return Err(fail(CliSource::Headless, CliErrorCode::InternalError, "--format requires a value", true)); + }; + if format != "json" { + return Err(fail(CliSource::Headless, CliErrorCode::InternalError, "Only --format json is supported", true)); + } + index += 2; + } else { + positionals.push(args[index].clone()); + index += 1; + } + } + + Ok(ParsedArgs { positionals }) +} + +async fn context(_args: &[String]) -> CliEnvelope { match crate::runtime_client::get_json("/context").await { Ok(data) => ok(CliSource::GuiRuntime, data), Err(_) => ok(CliSource::Headless, serde_json::json!({ "runtime": "headless" })), } } +async fn conn_list(_args: &[String]) -> CliEnvelope { + ok(CliSource::Headless, serde_json::json!({ "connections": [] })) +} + +async fn conn_show(_name: &str, _args: &[String]) -> CliEnvelope { + ok(CliSource::Headless, serde_json::json!({})) +} + +async fn schema_snapshot(args: &[String]) -> CliEnvelope { + let Some(_conn) = option_value(args, "--conn") else { + return fail(CliSource::Headless, CliErrorCode::ConnectionNotFound, "--conn is required", true); + }; + + fail(CliSource::Headless, CliErrorCode::InternalError, "Schema snapshot headless execution is not implemented", true) +} + +async fn safe_query(args: &[String]) -> CliEnvelope { + let Some(_conn) = option_value(args, "--conn") else { + return fail(CliSource::Headless, CliErrorCode::ConnectionNotFound, "--conn is required", true); + }; + let Some(_sql) = option_value(args, "--sql") else { + return fail(CliSource::Headless, CliErrorCode::QueryClassificationFailed, "--sql is required", true); + }; + + fail(CliSource::Headless, CliErrorCode::InternalError, "Safe query headless execution is not implemented", true) +} + +async fn handoff(args: &[String]) -> CliEnvelope { + let body = serde_json::json!({ + "conn": option_value(args, "--conn"), + "title": option_value(args, "--title"), + "sql": option_value(args, "--sql"), + }); + + match crate::runtime_client::post_json("/handoff", body).await { + Ok(data) => ok(CliSource::GuiRuntime, data), + Err(_) => fail(CliSource::Headless, CliErrorCode::InternalError, "Handoff payload is required", true), + } +} + +async fn selection(_args: &[String]) -> CliEnvelope { + match crate::runtime_client::get_json("/selection").await { + Ok(data) => ok(CliSource::GuiRuntime, data), + Err(_) => runtime_required("dbx selection requires DBX GUI runtime."), + } +} + +async fn result_current(args: &[String]) -> CliEnvelope { + let limit = option_value(args, "--limit").unwrap_or("50"); + + match crate::runtime_client::get_json(&format!("/result/current?limit={limit}")).await { + Ok(data) => ok(CliSource::GuiRuntime, data), + Err(_) => runtime_required("dbx result current requires DBX GUI runtime."), + } +} + fn runtime_required(message: &str) -> CliEnvelope { fail(CliSource::Headless, CliErrorCode::GuiRuntimeRequired, message, true) } + +fn option_value<'a>(args: &'a [String], key: &str) -> Option<&'a str> { + args.windows(2) + .find(|pair| pair[0] == key) + .map(|pair| pair[1].as_str()) +} + +#[cfg(test)] +mod tests { + use super::*; + use dbx_core::cli::CliErrorCode; + use std::sync::Mutex; + + static ENV_LOCK: Mutex<()> = Mutex::new(()); + + fn assert_failure_code(env: CliEnvelope, expected: CliErrorCode) { + match env { + CliEnvelope::Failure { error, .. } => assert_eq!(error.code, expected), + CliEnvelope::Success { .. } => panic!("expected failure envelope"), + } + } + + #[tokio::test] + async fn gui_only_commands_return_runtime_required_without_runtime() { + let _guard = ENV_LOCK.lock().unwrap(); + let dir = tempfile::tempdir().unwrap(); + std::env::set_var("DBX_APP_DATA_DIR", dir.path()); + + assert_failure_code( + dispatch(vec!["selection".into(), "--format".into(), "json".into()]).await, + CliErrorCode::GuiRuntimeRequired, + ); + assert_failure_code( + dispatch(vec![ + "result".into(), + "current".into(), + "--limit".into(), + "25".into(), + "--format".into(), + "json".into(), + ]) + .await, + CliErrorCode::GuiRuntimeRequired, + ); + + std::env::remove_var("DBX_APP_DATA_DIR"); + } + + #[tokio::test] + async fn recognizes_all_eight_cli_commands_with_json_format() { + let _guard = ENV_LOCK.lock().unwrap(); + let dir = tempfile::tempdir().unwrap(); + std::env::set_var("DBX_APP_DATA_DIR", dir.path()); + + let cases = [ + vec!["context", "--format", "json"], + vec!["conn", "list", "--format", "json"], + vec!["conn", "show", "__missing__", "--redacted", "--format", "json"], + vec!["schema", "snapshot", "--format", "json"], + vec!["safe-query", "--format", "json"], + vec!["handoff", "--format", "json"], + vec!["selection", "--format", "json"], + vec!["result", "current", "--limit", "50", "--format", "json"], + ]; + + for args in cases { + let env = dispatch(args.iter().map(|value| value.to_string()).collect()).await; + let json = serde_json::to_value(&env).unwrap(); + assert!(json.get("ok").is_some(), "missing ok for args: {args:?}"); + assert!(json.get("source").is_some(), "missing source for args: {args:?}"); + assert!( + json.get("data").is_some() || json.get("error").is_some(), + "missing data/error for args: {args:?}" + ); + } + + std::env::remove_var("DBX_APP_DATA_DIR"); + } + + #[tokio::test] + async fn unknown_command_returns_internal_error_envelope() { + assert_failure_code( + dispatch(vec!["not-a-command".into(), "--format".into(), "json".into()]).await, + CliErrorCode::InternalError, + ); + } + + #[tokio::test] + async fn rejects_non_json_format() { + assert_failure_code( + dispatch(vec!["context".into(), "--format".into(), "text".into()]).await, + CliErrorCode::InternalError, + ); + } +} diff --git a/crates/dbx-cli/src/runtime_client.rs b/crates/dbx-cli/src/runtime_client.rs index 940fd5399..45cb1898b 100644 --- a/crates/dbx-cli/src/runtime_client.rs +++ b/crates/dbx-cli/src/runtime_client.rs @@ -9,6 +9,10 @@ pub struct RuntimeDiscovery { } pub fn app_data_dir() -> PathBuf { + if let Ok(path) = std::env::var("DBX_APP_DATA_DIR") { + return PathBuf::from(path); + } + let home = std::env::var(if cfg!(windows) { "APPDATA" } else { "HOME" }) .unwrap_or_else(|_| ".".to_string()); @@ -31,12 +35,42 @@ pub async fn get_json(path: &str) -> Result { let runtime = load_runtime().ok_or_else(|| "runtime unavailable".to_string())?; let url = format!("http://127.0.0.1:{}{}", runtime.port, path); - reqwest::Client::new() + let response = reqwest::Client::new() .get(url) .bearer_auth(runtime.token) .send() .await - .map_err(|err| err.to_string())? + .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()) +} + +pub async fn post_json(path: &str, body: serde_json::Value) -> Result { + let runtime = load_runtime().ok_or_else(|| "runtime unavailable".to_string())?; + let url = format!("http://127.0.0.1:{}{}", runtime.port, path); + + let response = reqwest::Client::new() + .post(url) + .bearer_auth(runtime.token) + .json(&body) + .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()) diff --git a/crates/dbx-core/src/cli.rs b/crates/dbx-core/src/cli.rs index 26af58de2..dec743430 100644 --- a/crates/dbx-core/src/cli.rs +++ b/crates/dbx-core/src/cli.rs @@ -116,12 +116,8 @@ 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/storage.rs b/crates/dbx-core/src/storage.rs index 444792949..347caec01 100644 --- a/crates/dbx-core/src/storage.rs +++ b/crates/dbx-core/src/storage.rs @@ -144,11 +144,10 @@ 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(()); @@ -1058,8 +1057,10 @@ 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")); } From 37eeb8999c9064a011be7393f7f91e00f9de9fd7 Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 18:36:50 +0800 Subject: [PATCH 18/40] =?UTF-8?q?fix(dbx-cli):=20=E5=8A=A0=E5=9B=BA?= =?UTF-8?q?=E8=BF=90=E8=A1=8C=E6=97=B6=E5=91=BD=E4=BB=A4=E5=8F=82=E6=95=B0?= =?UTF-8?q?=E6=A0=A1=E9=AA=8C?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 严格校验未知 flag、缺失 option 值和 result current limit 范围 - 支持 handoff 的 SQL 来源校验及 headless 队列落盘 - 校验 runtime discovery 文件的 symlink、owner 和权限 --- Cargo.lock | 1 + crates/dbx-cli/Cargo.toml | 1 + crates/dbx-cli/src/commands.rs | 244 ++++++++++++++++++++++++--- crates/dbx-cli/src/main.rs | 5 +- crates/dbx-cli/src/runtime_client.rs | 149 ++++++++++++++-- 5 files changed, 350 insertions(+), 50 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 2b20b22c7..daf4f1613 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1679,6 +1679,7 @@ version = "0.5.2" dependencies = [ "chrono", "dbx-core", + "libc", "reqwest 0.12.28", "serde", "serde_json", diff --git a/crates/dbx-cli/Cargo.toml b/crates/dbx-cli/Cargo.toml index a81c00834..acb34da20 100644 --- a/crates/dbx-cli/Cargo.toml +++ b/crates/dbx-cli/Cargo.toml @@ -15,6 +15,7 @@ serde_json = "1.0" tokio = { version = "1", features = ["full"] } reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] } uuid = { version = "1", features = ["v4", "serde"] } +libc = "0.2" [dev-dependencies] tempfile = "3" diff --git a/crates/dbx-cli/src/commands.rs b/crates/dbx-cli/src/commands.rs index e97350dd6..a5d6c3f71 100644 --- a/crates/dbx-cli/src/commands.rs +++ b/crates/dbx-cli/src/commands.rs @@ -1,5 +1,8 @@ use dbx_core::cli::{fail, ok, CliEnvelope, CliErrorCode, CliSource}; +const DEFAULT_RESULT_LIMIT: u32 = 50; +const MAX_RESULT_LIMIT: u32 = 1000; + pub(crate) async fn run(args: Vec) -> Result<(), CliEnvelope<()>> { let output = dispatch(args).await; println!("{}", serde_json::to_string_pretty(&output).unwrap()); @@ -34,21 +37,44 @@ struct ParsedArgs { positionals: Vec, } +#[derive(Clone, Copy, PartialEq, Eq)] +enum FlagKind { + Value, + Bool, +} + fn parse_args(args: Vec) -> Result> { let mut positionals = Vec::new(); let mut index = 0; while index < args.len() { - if args[index] == "--format" { - let Some(format) = args.get(index + 1) else { - return Err(fail(CliSource::Headless, CliErrorCode::InternalError, "--format requires a value", true)); + let arg = &args[index]; + if arg.starts_with("--") { + let Some(kind) = flag_kind(&args, arg) else { + return Err(invalid_args(format!("Unknown flag: {arg}"))); }; - if format != "json" { - return Err(fail(CliSource::Headless, CliErrorCode::InternalError, "Only --format json is supported", true)); + + if kind == FlagKind::Value { + let Some(value) = args.get(index + 1) else { + return Err(invalid_args(format!("{arg} requires a value"))); + }; + if value.starts_with("--") { + return Err(invalid_args(format!("{arg} requires a value"))); + } + if arg == "--format" && value != "json" { + return Err(invalid_args("Only --format json is supported")); + } + if arg != "--format" { + positionals.push(arg.clone()); + positionals.push(value.clone()); + } + index += 2; + } else { + positionals.push(arg.clone()); + index += 1; } - index += 2; } else { - positionals.push(args[index].clone()); + positionals.push(arg.clone()); index += 1; } } @@ -56,6 +82,23 @@ fn parse_args(args: Vec) -> Result Option { + let first = args.first().map(String::as_str); + let second = args.get(1).map(String::as_str); + let value = FlagKind::Value; + let boolean = FlagKind::Bool; + + match (first, second, flag) { + (_, _, "--format") => Some(value), + (Some("conn"), Some("show"), "--redacted") => Some(boolean), + (Some("schema"), Some("snapshot"), "--conn" | "--db") => Some(value), + (Some("safe-query"), _, "--conn" | "--sql" | "--db") => Some(value), + (Some("handoff"), _, "--conn" | "--title" | "--sql" | "--sql-file" | "--description") => Some(value), + (Some("result"), Some("current"), "--limit") => Some(value), + _ => None, + } +} + async fn context(_args: &[String]) -> CliEnvelope { match crate::runtime_client::get_json("/context").await { Ok(data) => ok(CliSource::GuiRuntime, data), @@ -76,7 +119,12 @@ async fn schema_snapshot(args: &[String]) -> CliEnvelope { return fail(CliSource::Headless, CliErrorCode::ConnectionNotFound, "--conn is required", true); }; - fail(CliSource::Headless, CliErrorCode::InternalError, "Schema snapshot headless execution is not implemented", true) + fail( + CliSource::Headless, + CliErrorCode::InternalError, + "Schema snapshot headless execution is not implemented", + true, + ) } async fn safe_query(args: &[String]) -> CliEnvelope { @@ -91,15 +139,48 @@ async fn safe_query(args: &[String]) -> CliEnvelope { } async fn handoff(args: &[String]) -> CliEnvelope { - let body = serde_json::json!({ - "conn": option_value(args, "--conn"), - "title": option_value(args, "--title"), - "sql": option_value(args, "--sql"), - }); + let Some(conn) = required_option(args, "--conn") else { + return invalid_args("--conn is required"); + }; + let Some(title) = required_option(args, "--title") else { + return invalid_args("--title is required"); + }; + let sql_inline = option_value(args, "--sql"); + let sql_file = option_value(args, "--sql-file"); + let sql = match (sql_inline, sql_file) { + (Some(_), Some(_)) => return invalid_args("Use exactly one of --sql or --sql-file"), + (Some(sql), None) => sql.to_string(), + (None, Some(path)) => match std::fs::read_to_string(path) { + Ok(sql) => sql, + Err(err) => return invalid_args(format!("Failed to read --sql-file: {err}")), + }, + (None, None) => return invalid_args("Use exactly one of --sql or --sql-file"), + }; - match crate::runtime_client::post_json("/handoff", body).await { - Ok(data) => ok(CliSource::GuiRuntime, data), - Err(_) => fail(CliSource::Headless, CliErrorCode::InternalError, "Handoff payload is required", true), + if sql.trim().is_empty() { + return invalid_args("SQL must not be empty"); + } + + let risk = dbx_core::sql_safety::risk_for_connection(&sql, conn, None); + let item = dbx_core::handoff::HandoffItem::queued( + conn.to_string(), + conn.to_string(), + None, + title.to_string(), + option_value(args, "--description").map(str::to_string), + sql, + risk.operation_class, + risk.risk_level, + risk.is_production, + ); + + if let Ok(data) = crate::runtime_client::post_json("/handoff", serde_json::to_value(&item).unwrap()).await { + return ok(CliSource::GuiRuntime, data); + } + + match queue_handoff(&item).await { + Ok(()) => ok(CliSource::Headless, serde_json::json!({ "id": item.id, "status": "queued" })), + Err(err) => fail(CliSource::Headless, CliErrorCode::InternalError, err, false), } } @@ -111,9 +192,12 @@ async fn selection(_args: &[String]) -> CliEnvelope { } async fn result_current(args: &[String]) -> CliEnvelope { - let limit = option_value(args, "--limit").unwrap_or("50"); + let limit = match parse_result_limit(args) { + Ok(limit) => limit, + Err(err) => return err, + }; - match crate::runtime_client::get_json(&format!("/result/current?limit={limit}")).await { + match crate::runtime_client::get_json_with_query("/result/current", &[("limit", limit.to_string())]).await { Ok(data) => ok(CliSource::GuiRuntime, data), Err(_) => runtime_required("dbx result current requires DBX GUI runtime."), } @@ -123,19 +207,42 @@ fn runtime_required(message: &str) -> CliEnvelope { fail(CliSource::Headless, CliErrorCode::GuiRuntimeRequired, message, true) } +fn invalid_args(message: impl Into) -> CliEnvelope { + fail(CliSource::Headless, CliErrorCode::InternalError, message, true) +} + fn option_value<'a>(args: &'a [String], key: &str) -> Option<&'a str> { - args.windows(2) - .find(|pair| pair[0] == key) - .map(|pair| pair[1].as_str()) + args.windows(2).find(|pair| pair[0] == key).map(|pair| pair[1].as_str()) +} + +fn required_option<'a>(args: &'a [String], key: &str) -> Option<&'a str> { + option_value(args, key).map(str::trim).filter(|value| !value.is_empty()) +} + +fn parse_result_limit(args: &[String]) -> Result> { + let default_limit = DEFAULT_RESULT_LIMIT.to_string(); + let raw = option_value(args, "--limit").unwrap_or(&default_limit); + let limit = raw + .parse::() + .map_err(|_| invalid_args(format!("--limit must be a positive integer between 1 and {MAX_RESULT_LIMIT}")))?; + if !(1..=MAX_RESULT_LIMIT).contains(&limit) { + return Err(invalid_args(format!("--limit must be a positive integer between 1 and {MAX_RESULT_LIMIT}"))); + } + Ok(limit) +} + +async fn queue_handoff(item: &dbx_core::handoff::HandoffItem) -> Result<(), String> { + let app_dir = crate::runtime_client::app_data_dir(); + std::fs::create_dir_all(&app_dir).map_err(|err| err.to_string())?; + let storage = dbx_core::storage::Storage::open(&app_dir.join("dbx.db")).await?; + storage.save_handoff(item).await } #[cfg(test)] mod tests { use super::*; + use crate::runtime_client::ENV_LOCK; use dbx_core::cli::CliErrorCode; - use std::sync::Mutex; - - static ENV_LOCK: Mutex<()> = Mutex::new(()); fn assert_failure_code(env: CliEnvelope, expected: CliErrorCode) { match env { @@ -192,10 +299,7 @@ mod tests { let json = serde_json::to_value(&env).unwrap(); assert!(json.get("ok").is_some(), "missing ok for args: {args:?}"); assert!(json.get("source").is_some(), "missing source for args: {args:?}"); - assert!( - json.get("data").is_some() || json.get("error").is_some(), - "missing data/error for args: {args:?}" - ); + assert!(json.get("data").is_some() || json.get("error").is_some(), "missing data/error for args: {args:?}"); } std::env::remove_var("DBX_APP_DATA_DIR"); @@ -216,4 +320,88 @@ mod tests { CliErrorCode::InternalError, ); } + + #[tokio::test] + async fn rejects_unknown_flags_with_error_envelope() { + assert_failure_code( + dispatch(vec!["context".into(), "--unknown".into(), "value".into()]).await, + CliErrorCode::InternalError, + ); + } + + #[tokio::test] + async fn rejects_missing_option_values_with_error_envelope() { + assert_failure_code( + dispatch(vec!["safe-query".into(), "--conn".into(), "--sql".into(), "select 1".into()]).await, + CliErrorCode::InternalError, + ); + } + + #[tokio::test] + async fn validates_handoff_required_options_and_sql_source() { + assert_failure_code( + dispatch(vec!["handoff".into(), "--title".into(), "Review".into(), "--sql".into(), "select 1".into()]) + .await, + CliErrorCode::InternalError, + ); + assert_failure_code( + dispatch(vec!["handoff".into(), "--conn".into(), "local".into(), "--sql".into(), "select 1".into()]).await, + CliErrorCode::InternalError, + ); + assert_failure_code( + dispatch(vec![ + "handoff".into(), + "--conn".into(), + "local".into(), + "--title".into(), + "Review".into(), + "--sql".into(), + "select 1".into(), + "--sql-file".into(), + "query.sql".into(), + ]) + .await, + CliErrorCode::InternalError, + ); + } + + #[tokio::test] + async fn handoff_queues_sql_file_when_runtime_is_unavailable() { + let _guard = ENV_LOCK.lock().unwrap(); + let dir = tempfile::tempdir().unwrap(); + 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()); + + let env = dispatch(vec![ + "handoff".into(), + "--conn".into(), + "local".into(), + "--title".into(), + "Review".into(), + "--sql-file".into(), + sql_file.display().to_string(), + ]) + .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:?}"), + } + + 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"] { + assert_failure_code( + dispatch(vec!["result".into(), "current".into(), "--limit".into(), limit.into()]).await, + CliErrorCode::InternalError, + ); + } + } } diff --git a/crates/dbx-cli/src/main.rs b/crates/dbx-cli/src/main.rs index f34308a75..9a5a19356 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 45cb1898b..effa954a0 100644 --- a/crates/dbx-cli/src/runtime_client.rs +++ b/crates/dbx-cli/src/runtime_client.rs @@ -1,6 +1,10 @@ use serde::{Deserialize, Serialize}; +use std::fs::Metadata; use std::path::PathBuf; +#[cfg(test)] +pub(crate) static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct RuntimeDiscovery { @@ -13,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") @@ -27,35 +30,47 @@ pub fn app_data_dir() -> PathBuf { pub fn load_runtime() -> Option { let path = app_data_dir().join("agent-runtime.json"); + let metadata = std::fs::symlink_metadata(&path).ok()?; + if !is_secure_runtime_file(&metadata) { + return None; + } let json = std::fs::read_to_string(path).ok()?; serde_json::from_str(&json).ok() } pub async fn get_json(path: &str) -> Result { let runtime = load_runtime().ok_or_else(|| "runtime unavailable".to_string())?; - let url = format!("http://127.0.0.1:{}{}", runtime.port, path); + 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 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()) } pub async fn post_json(path: &str, body: serde_json::Value) -> Result { let runtime = load_runtime().ok_or_else(|| "runtime unavailable".to_string())?; - let url = format!("http://127.0.0.1:{}{}", runtime.port, path); + let url = runtime_url(&runtime, path, &[])?; let response = reqwest::Client::new() .post(url) @@ -70,8 +85,106 @@ pub async fn post_json(path: &str, body: serde_json::Value) -> Result Result { + if path.contains('\r') || path.contains('\n') || path.contains("://") { + return Err("invalid runtime path".to_string()); + } + + let mut url = reqwest::Url::parse(&format!("http://127.0.0.1:{}/", runtime.port)).map_err(|err| err.to_string())?; + url.set_path(path.trim_start_matches('/')); + { + let mut pairs = url.query_pairs_mut(); + for (key, value) in query { + pairs.append_pair(key, value); + } + } + Ok(url) +} + +fn is_secure_runtime_file(metadata: &Metadata) -> bool { + if metadata.file_type().is_symlink() || !metadata.is_file() { + return false; + } + runtime_file_owner_and_mode_are_secure(metadata) +} + +#[cfg(unix)] +fn runtime_file_owner_and_mode_are_secure(metadata: &Metadata) -> bool { + use std::os::unix::fs::MetadataExt; + + let owner_only = metadata.mode() & 0o077 == 0; + let owned_by_effective_user = metadata.uid() == unsafe { libc::geteuid() }; + owner_only && owned_by_effective_user +} + +#[cfg(not(unix))] +fn runtime_file_owner_and_mode_are_secure(_metadata: &Metadata) -> bool { + true +} + +#[cfg(test)] +mod tests { + use super::*; + + #[cfg(unix)] + fn set_mode(path: &std::path::Path, mode: u32) { + use std::os::unix::fs::PermissionsExt; + + let mut permissions = std::fs::metadata(path).unwrap().permissions(); + permissions.set_mode(mode); + std::fs::set_permissions(path, permissions).unwrap(); + } + + #[test] + fn load_runtime_rejects_symlink_discovery_file() { + let _guard = ENV_LOCK.lock().unwrap(); + let dir = tempfile::tempdir().unwrap(); + let target = dir.path().join("target.json"); + let link = dir.path().join("agent-runtime.json"); + std::fs::write(&target, r#"{"port":4321,"token":"secret"}"#).unwrap(); + #[cfg(unix)] + std::os::unix::fs::symlink(&target, &link).unwrap(); + #[cfg(not(unix))] + std::fs::write(&link, r#"{"port":4321,"token":"secret"}"#).unwrap(); + std::env::set_var("DBX_APP_DATA_DIR", dir.path()); + + assert!(load_runtime().is_none()); + + std::env::remove_var("DBX_APP_DATA_DIR"); + } + + #[cfg(unix)] + #[test] + fn load_runtime_rejects_group_or_world_accessible_discovery_file() { + let _guard = ENV_LOCK.lock().unwrap(); + let dir = tempfile::tempdir().unwrap(); + let discovery = dir.path().join("agent-runtime.json"); + std::fs::write(&discovery, r#"{"port":4321,"token":"secret"}"#).unwrap(); + set_mode(&discovery, 0o644); + std::env::set_var("DBX_APP_DATA_DIR", dir.path()); + + assert!(load_runtime().is_none()); + + std::env::remove_var("DBX_APP_DATA_DIR"); + } + + #[cfg(unix)] + #[test] + fn load_runtime_accepts_owner_only_discovery_file() { + let _guard = ENV_LOCK.lock().unwrap(); + let dir = tempfile::tempdir().unwrap(); + let discovery = dir.path().join("agent-runtime.json"); + std::fs::write(&discovery, r#"{"port":4321,"token":"secret"}"#).unwrap(); + set_mode(&discovery, 0o600); + std::env::set_var("DBX_APP_DATA_DIR", dir.path()); + + let runtime = load_runtime().expect("secure runtime discovery should load"); + assert_eq!(runtime.port, 4321); + assert_eq!(runtime.token, "secret"); + + std::env::remove_var("DBX_APP_DATA_DIR"); + } } From dc13dd33762b559cbab5a3f3571938c20109d743 Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 18:42:39 +0800 Subject: [PATCH 19/40] =?UTF-8?q?fix(cli):=20=E4=BF=AE=E5=A4=8D=E5=91=BD?= =?UTF-8?q?=E4=BB=A4=E5=8F=82=E6=95=B0=E8=A7=A3=E6=9E=90?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 允许自由文本选项值以双短横开头 - 支持全局 --format json 出现在命令前、中、后 - 拒绝各命令多余位置参数并补充回归测试 --- crates/dbx-cli/src/commands.rs | 197 ++++++++++++++++++++++++++++----- 1 file changed, 172 insertions(+), 25 deletions(-) diff --git a/crates/dbx-cli/src/commands.rs b/crates/dbx-cli/src/commands.rs index a5d6c3f71..db4b71422 100644 --- a/crates/dbx-cli/src/commands.rs +++ b/crates/dbx-cli/src/commands.rs @@ -43,31 +43,42 @@ enum FlagKind { Bool, } +#[derive(Clone, Copy)] +struct FlagSpec { + kind: FlagKind, + allow_dash_value: bool, +} + fn parse_args(args: Vec) -> Result> { let mut positionals = Vec::new(); let mut index = 0; while index < args.len() { let arg = &args[index]; - if arg.starts_with("--") { - let Some(kind) = flag_kind(&args, arg) else { + if arg == "--format" { + let Some(value) = args.get(index + 1) else { + return Err(invalid_args(format!("{arg} requires a value"))); + }; + if value != "json" { + return Err(invalid_args("Only --format json is supported")); + } + index += 2; + } else if arg.starts_with("--") { + let Some(spec) = + flag_spec(positionals.first().map(String::as_str), positionals.get(1).map(String::as_str), arg) + else { return Err(invalid_args(format!("Unknown flag: {arg}"))); }; - if kind == FlagKind::Value { + if spec.kind == FlagKind::Value { let Some(value) = args.get(index + 1) else { return Err(invalid_args(format!("{arg} requires a value"))); }; - if value.starts_with("--") { + if !spec.allow_dash_value && value.starts_with("--") { return Err(invalid_args(format!("{arg} requires a value"))); } - if arg == "--format" && value != "json" { - return Err(invalid_args("Only --format json is supported")); - } - if arg != "--format" { - positionals.push(arg.clone()); - positionals.push(value.clone()); - } + positionals.push(arg.clone()); + positionals.push(value.clone()); index += 2; } else { positionals.push(arg.clone()); @@ -82,39 +93,77 @@ fn parse_args(args: Vec) -> Result Option { - let first = args.first().map(String::as_str); - let second = args.get(1).map(String::as_str); +fn flag_spec(command: Option<&str>, subcommand: Option<&str>, flag: &str) -> Option { let value = FlagKind::Value; let boolean = FlagKind::Bool; + let normal_value = FlagSpec { kind: value, allow_dash_value: false }; + let free_text_value = FlagSpec { kind: value, allow_dash_value: true }; + let boolean_flag = FlagSpec { kind: boolean, allow_dash_value: false }; - match (first, second, flag) { - (_, _, "--format") => Some(value), - (Some("conn"), Some("show"), "--redacted") => Some(boolean), - (Some("schema"), Some("snapshot"), "--conn" | "--db") => Some(value), - (Some("safe-query"), _, "--conn" | "--sql" | "--db") => Some(value), - (Some("handoff"), _, "--conn" | "--title" | "--sql" | "--sql-file" | "--description") => Some(value), - (Some("result"), Some("current"), "--limit") => Some(value), + match (command, subcommand, flag) { + (Some("conn"), Some("show"), "--redacted") => Some(boolean_flag), + (Some("schema"), Some("snapshot"), "--conn" | "--db") => Some(normal_value), + (Some("safe-query"), _, "--conn" | "--db") => Some(normal_value), + (Some("safe-query"), _, "--sql") => Some(free_text_value), + (Some("handoff"), _, "--conn" | "--sql-file") => Some(normal_value), + (Some("handoff"), _, "--title" | "--sql" | "--description") => Some(free_text_value), + (Some("result"), Some("current"), "--limit") => Some(normal_value), _ => None, } } -async fn context(_args: &[String]) -> CliEnvelope { +fn reject_unexpected_positionals( + args: &[String], + command: &str, + subcommand: Option<&str>, +) -> Result<(), CliEnvelope> { + let mut index = 0; + while index < args.len() { + let arg = &args[index]; + let Some(spec) = flag_spec(Some(command), subcommand, arg) else { + return Err(invalid_args(format!("Unexpected positional argument: {arg}"))); + }; + + index += match spec.kind { + FlagKind::Value => 2, + FlagKind::Bool => 1, + }; + } + Ok(()) +} + +async fn context(args: &[String]) -> CliEnvelope { + if let Err(err) = reject_unexpected_positionals(args, "context", None) { + return err; + } + match crate::runtime_client::get_json("/context").await { Ok(data) => ok(CliSource::GuiRuntime, data), Err(_) => ok(CliSource::Headless, serde_json::json!({ "runtime": "headless" })), } } -async fn conn_list(_args: &[String]) -> CliEnvelope { +async fn conn_list(args: &[String]) -> CliEnvelope { + if let Err(err) = reject_unexpected_positionals(args, "conn", Some("list")) { + return err; + } + ok(CliSource::Headless, serde_json::json!({ "connections": [] })) } -async fn conn_show(_name: &str, _args: &[String]) -> CliEnvelope { +async fn conn_show(_name: &str, args: &[String]) -> CliEnvelope { + if let Err(err) = reject_unexpected_positionals(args, "conn", Some("show")) { + return err; + } + ok(CliSource::Headless, serde_json::json!({})) } async fn schema_snapshot(args: &[String]) -> CliEnvelope { + if let Err(err) = reject_unexpected_positionals(args, "schema", Some("snapshot")) { + return err; + } + let Some(_conn) = option_value(args, "--conn") else { return fail(CliSource::Headless, CliErrorCode::ConnectionNotFound, "--conn is required", true); }; @@ -128,6 +177,10 @@ async fn schema_snapshot(args: &[String]) -> CliEnvelope { } async fn safe_query(args: &[String]) -> CliEnvelope { + if let Err(err) = reject_unexpected_positionals(args, "safe-query", None) { + return err; + } + let Some(_conn) = option_value(args, "--conn") else { return fail(CliSource::Headless, CliErrorCode::ConnectionNotFound, "--conn is required", true); }; @@ -139,6 +192,10 @@ async fn safe_query(args: &[String]) -> CliEnvelope { } async fn handoff(args: &[String]) -> CliEnvelope { + if let Err(err) = reject_unexpected_positionals(args, "handoff", None) { + return err; + } + let Some(conn) = required_option(args, "--conn") else { return invalid_args("--conn is required"); }; @@ -184,7 +241,11 @@ async fn handoff(args: &[String]) -> CliEnvelope { } } -async fn selection(_args: &[String]) -> CliEnvelope { +async fn selection(args: &[String]) -> CliEnvelope { + if let Err(err) = reject_unexpected_positionals(args, "selection", None) { + return err; + } + match crate::runtime_client::get_json("/selection").await { Ok(data) => ok(CliSource::GuiRuntime, data), Err(_) => runtime_required("dbx selection requires DBX GUI runtime."), @@ -192,6 +253,10 @@ async fn selection(_args: &[String]) -> CliEnvelope { } async fn result_current(args: &[String]) -> CliEnvelope { + if let Err(err) = reject_unexpected_positionals(args, "result", Some("current")) { + return err; + } + let limit = match parse_result_limit(args) { Ok(limit) => limit, Err(err) => return err, @@ -251,6 +316,88 @@ mod tests { } } + fn assert_failure_message_contains(env: CliEnvelope, expected: &str) { + match env { + CliEnvelope::Failure { error, .. } => assert!( + error.message.contains(expected), + "expected error message to contain {expected:?}, got {:?}", + error.message + ), + CliEnvelope::Success { .. } => panic!("expected failure envelope"), + } + } + + #[test] + fn parser_allows_free_text_option_values_to_start_with_dashes() { + let parsed = parse_args(vec![ + "handoff".into(), + "--conn".into(), + "local".into(), + "--title".into(), + "-- review generated SQL".into(), + "--sql".into(), + "-- explain select 1".into(), + "--description".into(), + "-- optional note".into(), + ]) + .expect("free-text values beginning with -- should parse"); + + assert_eq!( + parsed.positionals, + vec![ + "handoff", + "--conn", + "local", + "--title", + "-- review generated SQL", + "--sql", + "-- explain select 1", + "--description", + "-- optional note", + ] + ); + } + + #[test] + fn parser_accepts_global_json_format_before_between_and_after_command_args() { + let cases = [ + vec!["--format", "json", "safe-query", "--conn", "local", "--sql", "select 1"], + vec!["safe-query", "--conn", "local", "--format", "json", "--sql", "select 1"], + vec!["safe-query", "--conn", "local", "--sql", "select 1", "--format", "json"], + ]; + + for args in cases { + let parsed = parse_args(args.iter().map(|value| value.to_string()).collect()) + .unwrap_or_else(|err| panic!("expected parse success for {args:?}, got {err:?}")); + assert_eq!(parsed.positionals, vec!["safe-query", "--conn", "local", "--sql", "select 1"]); + } + } + + #[tokio::test] + async fn rejects_extra_positionals_for_each_command() { + let _guard = ENV_LOCK.lock().unwrap(); + let dir = tempfile::tempdir().unwrap(); + std::env::set_var("DBX_APP_DATA_DIR", dir.path()); + + let cases = [ + vec!["context", "extra"], + vec!["conn", "list", "extra"], + vec!["conn", "show", "local", "extra"], + vec!["schema", "snapshot", "--conn", "local", "extra"], + vec!["safe-query", "--conn", "local", "--sql", "select 1", "extra"], + vec!["handoff", "--conn", "local", "--title", "Review", "--sql", "select 1", "extra"], + vec!["selection", "extra"], + vec!["result", "current", "--limit", "50", "extra"], + ]; + + for args in cases { + let env = dispatch(args.iter().map(|value| value.to_string()).collect()).await; + assert_failure_message_contains(env, "Unexpected positional argument"); + } + + std::env::remove_var("DBX_APP_DATA_DIR"); + } + #[tokio::test] async fn gui_only_commands_return_runtime_required_without_runtime() { let _guard = ENV_LOCK.lock().unwrap(); From 9b6d058cd5ceea10f0749f33ac9c1644f35f5cb2 Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 18:47:43 +0800 Subject: [PATCH 20/40] =?UTF-8?q?feat(dbx-cli):=20=E5=AE=9E=E7=8E=B0=20hea?= =?UTF-8?q?dless=20=E8=BF=9E=E6=8E=A5=E8=AF=BB=E5=8F=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 从 DBX Storage 读取 context、conn list 和 conn show 的连接配置 - 输出脱敏连接 DTO 与风险元数据，避免泄露密码和 SSH secret - 覆盖连接不存在、名称歧义和 fixture 存储读取测试 --- crates/dbx-cli/src/commands.rs | 240 ++++++++++++++++++++++++++++++++- 1 file changed, 236 insertions(+), 4 deletions(-) diff --git a/crates/dbx-cli/src/commands.rs b/crates/dbx-cli/src/commands.rs index db4b71422..fda97e9eb 100644 --- a/crates/dbx-cli/src/commands.rs +++ b/crates/dbx-cli/src/commands.rs @@ -132,6 +132,71 @@ fn reject_unexpected_positionals( Ok(()) } +async fn open_state() -> Result { + let app_dir = crate::runtime_client::app_data_dir(); + std::fs::create_dir_all(&app_dir).map_err(|err| err.to_string())?; + let storage = dbx_core::storage::Storage::open(&app_dir.join("dbx.db")).await?; + Ok(dbx_core::connection::AppState::new(storage)) +} + +fn redacted_config(config: &dbx_core::models::connection::ConnectionConfig) -> serde_json::Value { + let risk = dbx_core::sql_safety::risk_for_connection("SELECT 1", &config.name, config.color.as_deref()); + serde_json::json!({ + "id": config.id, + "name": config.name, + "databaseType": config.db_type, + "driverProfile": config.driver_profile, + "driverLabel": config.driver_label, + "defaultDatabase": config.database, + "color": config.color, + "sshEnabled": config.ssh_enabled, + "redactedUrl": redacted_url(config), + "risk": risk, + }) +} + +fn redacted_url(config: &dbx_core::models::connection::ConnectionConfig) -> String { + redact_uri_credentials(&config.redacted_connection_url()) +} + +fn redact_uri_credentials(value: &str) -> String { + let Some(scheme_index) = value.find("://") else { + return value.to_string(); + }; + let authority_start = scheme_index + 3; + let rest = &value[authority_start..]; + let authority_len = rest.find(['/', '?', '#']).unwrap_or(rest.len()); + let authority = &rest[..authority_len]; + let Some(at_index) = authority.rfind('@') else { + return value.to_string(); + }; + + format!("{}{}{}", &value[..authority_start], &authority[at_index + 1..], &rest[authority_len..]) +} + +async fn load_headless_connections( +) -> Result, CliEnvelope> { + let state = open_state().await.map_err(|err| fail(CliSource::Headless, CliErrorCode::InternalError, err, false))?; + state + .storage + .load_connections() + .await + .map_err(|err| fail(CliSource::Headless, CliErrorCode::InternalError, err, false)) +} + +async fn find_connection( + name: &str, +) -> Result> { + let configs = load_headless_connections().await?; + let matches: Vec<_> = configs.into_iter().filter(|config| config.name == name || config.id == name).collect(); + + match matches.len() { + 1 => Ok(matches.into_iter().next().unwrap()), + 0 => Err(fail(CliSource::Headless, CliErrorCode::ConnectionNotFound, "Connection not found", true)), + _ => Err(fail(CliSource::Headless, CliErrorCode::AmbiguousConnection, "Connection name is ambiguous", true)), + } +} + async fn context(args: &[String]) -> CliEnvelope { if let Err(err) = reject_unexpected_positionals(args, "context", None) { return err; @@ -139,7 +204,20 @@ async fn context(args: &[String]) -> CliEnvelope { match crate::runtime_client::get_json("/context").await { Ok(data) => ok(CliSource::GuiRuntime, data), - Err(_) => ok(CliSource::Headless, serde_json::json!({ "runtime": "headless" })), + Err(_) => { + let configs = match load_headless_connections().await { + Ok(configs) => configs, + Err(err) => return err, + }; + ok( + CliSource::Headless, + serde_json::json!({ + "runtime": "headless", + "activeConnection": configs.first().map(redacted_config), + "configSource": crate::runtime_client::app_data_dir().join("dbx.db").display().to_string(), + }), + ) + } } } @@ -148,15 +226,26 @@ async fn conn_list(args: &[String]) -> CliEnvelope { return err; } - ok(CliSource::Headless, serde_json::json!({ "connections": [] })) + match load_headless_connections().await { + Ok(configs) => ok( + CliSource::Headless, + serde_json::json!({ + "connections": configs.iter().map(redacted_config).collect::>(), + }), + ), + Err(err) => err, + } } -async fn conn_show(_name: &str, args: &[String]) -> CliEnvelope { +async fn conn_show(name: &str, args: &[String]) -> CliEnvelope { if let Err(err) = reject_unexpected_positionals(args, "conn", Some("show")) { return err; } - ok(CliSource::Headless, serde_json::json!({})) + match find_connection(name).await { + Ok(config) => ok(CliSource::Headless, redacted_config(&config)), + Err(err) => err, + } } async fn schema_snapshot(args: &[String]) -> CliEnvelope { @@ -308,6 +397,7 @@ mod tests { use super::*; use crate::runtime_client::ENV_LOCK; use dbx_core::cli::CliErrorCode; + use dbx_core::models::connection::{ConnectionConfig, DatabaseType}; fn assert_failure_code(env: CliEnvelope, expected: CliErrorCode) { match env { @@ -327,6 +417,53 @@ mod tests { } } + fn mysql_fixture(id: &str, name: &str, color: Option<&str>, password: &str) -> ConnectionConfig { + ConnectionConfig { + id: id.to_string(), + name: name.to_string(), + db_type: DatabaseType::Mysql, + driver_profile: None, + driver_label: Some("MySQL".to_string()), + url_params: None, + host: "127.0.0.1".to_string(), + port: 3306, + username: "root".to_string(), + password: password.to_string(), + database: Some("app".to_string()), + color: color.map(str::to_string), + ssh_enabled: true, + ssh_host: "bastion.internal".to_string(), + ssh_port: 22, + ssh_user: "deploy".to_string(), + ssh_password: "ssh-secret".to_string(), + ssh_key_path: String::new(), + ssh_key_passphrase: "key-secret".to_string(), + ssh_expose_lan: false, + ssh_connect_timeout_secs: dbx_core::models::connection::default_ssh_connect_timeout_secs(), + ssl: false, + sysdba: false, + connection_string: Some(format!("mysql://root:{password}@127.0.0.1:3306/app")), + jdbc_driver_class: None, + jdbc_driver_paths: Vec::new(), + } + } + + async fn seed_connections(dir: &std::path::Path, configs: &[ConnectionConfig]) { + std::fs::create_dir_all(dir).unwrap(); + let storage = dbx_core::storage::Storage::open(&dir.join("dbx.db")).await.unwrap(); + storage.save_connections(configs).await.unwrap(); + } + + fn success_data(env: CliEnvelope) -> serde_json::Value { + match env { + CliEnvelope::Success { source, data, .. } => { + assert_eq!(source, CliSource::Headless); + data + } + CliEnvelope::Failure { error, .. } => panic!("expected success envelope, got {error:?}"), + } + } + #[test] fn parser_allows_free_text_option_values_to_start_with_dashes() { let parsed = parse_args(vec![ @@ -452,6 +589,101 @@ mod tests { std::env::remove_var("DBX_APP_DATA_DIR"); } + #[tokio::test] + async fn context_reads_headless_storage_and_returns_redacted_active_connection() { + 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"), "super-secret")]).await; + + let data = success_data(dispatch(vec!["context".into(), "--format".into(), "json".into()]).await); + + assert_eq!(data["runtime"], "headless"); + assert_eq!(data["activeConnection"]["id"], "prod-id"); + assert_eq!(data["activeConnection"]["name"], "prod-main"); + assert_eq!(data["activeConnection"]["risk"]["isProduction"], true); + let json = serde_json::to_string(&data).unwrap(); + assert!(json.contains("configSource")); + assert!(!json.contains("super-secret")); + assert!(!json.contains("ssh-secret")); + assert!(!json.contains("key-secret")); + + std::env::remove_var("DBX_APP_DATA_DIR"); + } + + #[tokio::test] + async fn conn_list_reads_storage_and_returns_redacted_connection_dtos_with_risk_metadata() { + 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"), + mysql_fixture("dev-id", "dev-main", Some("#22c55e"), "dev-secret"), + ], + ) + .await; + + let data = success_data(dispatch(vec!["conn".into(), "list".into(), "--format".into(), "json".into()]).await); + let connections = data["connections"].as_array().expect("connections should be an array"); + + assert_eq!(connections.len(), 2); + assert_eq!(connections[0]["id"], "prod-id"); + assert_eq!(connections[0]["redactedUrl"], "mysql://127.0.0.1:3306/app?ssl-mode=preferred&charset=utf8mb4"); + assert_eq!(connections[0]["risk"]["isProduction"], true); + assert_eq!(connections[0]["risk"]["productionReason"], "red connection color"); + assert_eq!(connections[1]["risk"]["isProduction"], false); + let json = serde_json::to_string(&data).unwrap(); + assert!(!json.contains("prod-secret")); + assert!(!json.contains("dev-secret")); + assert!(!json.contains("root:")); + + std::env::remove_var("DBX_APP_DATA_DIR"); + } + + #[tokio::test] + async fn conn_show_supports_id_lookup_and_returns_not_found_or_ambiguous_errors() { + 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"), + mysql_fixture("dup-1", "shared", None, "first-secret"), + mysql_fixture("dup-2", "shared", None, "second-secret"), + ], + ) + .await; + + let shown = success_data( + dispatch(vec![ + "conn".into(), + "show".into(), + "prod-id".into(), + "--redacted".into(), + "--format".into(), + "json".into(), + ]) + .await, + ); + assert_eq!(shown["id"], "prod-id"); + assert_eq!(shown["risk"]["isProduction"], true); + assert!(!serde_json::to_string(&shown).unwrap().contains("prod-secret")); + + assert_failure_code( + dispatch(vec!["conn".into(), "show".into(), "__missing__".into(), "--format".into(), "json".into()]).await, + CliErrorCode::ConnectionNotFound, + ); + assert_failure_code( + dispatch(vec!["conn".into(), "show".into(), "shared".into(), "--format".into(), "json".into()]).await, + CliErrorCode::AmbiguousConnection, + ); + + std::env::remove_var("DBX_APP_DATA_DIR"); + } + #[tokio::test] async fn unknown_command_returns_internal_error_envelope() { assert_failure_code( From 4c84cd0559cf4d4600cf4014178387b7c585aa68 Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 18:52:04 +0800 Subject: [PATCH 21/40] =?UTF-8?q?fix(cli):=20=E4=BF=AE=E5=A4=8D=E8=BF=9E?= =?UTF-8?q?=E6=8E=A5=E5=B1=95=E7=A4=BA=E7=9A=84=E5=8C=B9=E9=85=8D=E4=B8=8E?= =?UTF-8?q?=E8=84=B1=E6=95=8F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 优先按唯一连接 ID 命中 conn show，避免被同名连接歧义阻断 - headless context 不再输出绝对 dbx.db 路径 - 脱敏 SQLite 和 DuckDB redactedUrl 中的本地文件路径 --- crates/dbx-cli/src/commands.rs | 102 ++++++++++++++++++++++++++++++++- 1 file changed, 100 insertions(+), 2 deletions(-) diff --git a/crates/dbx-cli/src/commands.rs b/crates/dbx-cli/src/commands.rs index fda97e9eb..dbf54b30a 100644 --- a/crates/dbx-cli/src/commands.rs +++ b/crates/dbx-cli/src/commands.rs @@ -156,9 +156,21 @@ fn redacted_config(config: &dbx_core::models::connection::ConnectionConfig) -> s } fn redacted_url(config: &dbx_core::models::connection::ConnectionConfig) -> String { + if matches!( + config.db_type, + dbx_core::models::connection::DatabaseType::Sqlite | dbx_core::models::connection::DatabaseType::DuckDb + ) { + return redact_embedded_path_url(&config.redacted_connection_url()); + } + redact_uri_credentials(&config.redacted_connection_url()) } +fn redact_embedded_path_url(value: &str) -> String { + let suffix_start = value.find(['?', '#']).unwrap_or(value.len()); + format!("{}", &value[suffix_start..]) +} + fn redact_uri_credentials(value: &str) -> String { let Some(scheme_index) = value.find("://") else { return value.to_string(); @@ -188,7 +200,21 @@ async fn find_connection( name: &str, ) -> Result> { let configs = load_headless_connections().await?; - let matches: Vec<_> = configs.into_iter().filter(|config| config.name == name || config.id == name).collect(); + let id_matches: Vec<_> = configs.iter().filter(|config| config.id == name).collect(); + match id_matches.len() { + 1 => return Ok(id_matches[0].clone()), + 0 => {} + _ => { + return Err(fail( + CliSource::Headless, + CliErrorCode::AmbiguousConnection, + "Connection id is ambiguous", + true, + )); + } + } + + let matches: Vec<_> = configs.into_iter().filter(|config| config.name == name).collect(); match matches.len() { 1 => Ok(matches.into_iter().next().unwrap()), @@ -214,7 +240,7 @@ async fn context(args: &[String]) -> CliEnvelope { serde_json::json!({ "runtime": "headless", "activeConnection": configs.first().map(redacted_config), - "configSource": crate::runtime_client::app_data_dir().join("dbx.db").display().to_string(), + "configSource": "headless", }), ) } @@ -448,6 +474,18 @@ mod tests { } } + fn embedded_fixture(id: &str, name: &str, db_type: DatabaseType, path: &std::path::Path) -> ConnectionConfig { + let mut config = mysql_fixture(id, name, None, ""); + config.db_type = db_type; + config.host = path.display().to_string(); + config.port = 0; + config.username = String::new(); + config.password = String::new(); + config.database = None; + config.connection_string = None; + config + } + async fn seed_connections(dir: &std::path::Path, configs: &[ConnectionConfig]) { std::fs::create_dir_all(dir).unwrap(); let storage = dbx_core::storage::Storage::open(&dir.join("dbx.db")).await.unwrap(); @@ -602,8 +640,10 @@ mod tests { assert_eq!(data["activeConnection"]["id"], "prod-id"); assert_eq!(data["activeConnection"]["name"], "prod-main"); assert_eq!(data["activeConnection"]["risk"]["isProduction"], true); + assert_eq!(data["configSource"], "headless"); let json = serde_json::to_string(&data).unwrap(); assert!(json.contains("configSource")); + assert!(!json.contains(&dir.path().join("dbx.db").display().to_string())); assert!(!json.contains("super-secret")); assert!(!json.contains("ssh-secret")); assert!(!json.contains("key-secret")); @@ -642,6 +682,35 @@ mod tests { std::env::remove_var("DBX_APP_DATA_DIR"); } + #[tokio::test] + async fn conn_list_redacts_sqlite_and_duckdb_file_paths() { + let _guard = ENV_LOCK.lock().unwrap(); + let dir = tempfile::tempdir().unwrap(); + std::env::set_var("DBX_APP_DATA_DIR", dir.path()); + let sqlite_path = dir.path().join("private").join("app.sqlite"); + let duckdb_path = dir.path().join("private").join("warehouse.duckdb"); + seed_connections( + dir.path(), + &[ + embedded_fixture("sqlite-id", "local-sqlite", DatabaseType::Sqlite, &sqlite_path), + embedded_fixture("duckdb-id", "local-duckdb", DatabaseType::DuckDb, &duckdb_path), + ], + ) + .await; + + let data = success_data(dispatch(vec!["conn".into(), "list".into(), "--format".into(), "json".into()]).await); + let connections = data["connections"].as_array().expect("connections should be an array"); + + assert_eq!(connections[0]["redactedUrl"], "?mode=rwc"); + assert_eq!(connections[1]["redactedUrl"], "?mode=rwc"); + let json = serde_json::to_string(&data).unwrap(); + assert!(!json.contains("app.sqlite")); + assert!(!json.contains("warehouse.duckdb")); + assert!(!json.contains(dir.path().to_str().unwrap())); + + std::env::remove_var("DBX_APP_DATA_DIR"); + } + #[tokio::test] async fn conn_show_supports_id_lookup_and_returns_not_found_or_ambiguous_errors() { let _guard = ENV_LOCK.lock().unwrap(); @@ -684,6 +753,35 @@ mod tests { std::env::remove_var("DBX_APP_DATA_DIR"); } + #[tokio::test] + async fn conn_show_unique_id_match_takes_precedence_over_ambiguous_name() { + 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("shared", "id-target", Some("#22c55e"), "target-secret"), + mysql_fixture("dup-1", "shared", None, "first-secret"), + mysql_fixture("dup-2", "shared", None, "second-secret"), + ], + ) + .await; + + let shown = success_data( + dispatch(vec!["conn".into(), "show".into(), "shared".into(), "--format".into(), "json".into()]).await, + ); + + assert_eq!(shown["id"], "shared"); + assert_eq!(shown["name"], "id-target"); + let json = serde_json::to_string(&shown).unwrap(); + assert!(!json.contains("target-secret")); + assert!(!json.contains("first-secret")); + assert!(!json.contains("second-secret")); + + std::env::remove_var("DBX_APP_DATA_DIR"); + } + #[tokio::test] async fn unknown_command_returns_internal_error_envelope() { assert_failure_code( From 7486e2c7b3a74c2cdd1944dc6079dc58820228b1 Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 18:59:00 +0800 Subject: [PATCH 22/40] =?UTF-8?q?feat(dbx-cli):=20=E5=AE=9E=E7=8E=B0=20hea?= =?UTF-8?q?dless=20schema=20=E4=B8=8E=20safe-query?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 调用 dbx_core::schema_snapshot 生成 headless schema snapshot - safe-query 从 storage 加载连接并按风险阻断非 READ SQL - 补充 SQLite fixture 覆盖 READ、limit 与阻断错误 --- crates/dbx-cli/src/commands.rs | 250 +++++++++++++++++++++++++++++++-- 1 file changed, 239 insertions(+), 11 deletions(-) diff --git a/crates/dbx-cli/src/commands.rs b/crates/dbx-cli/src/commands.rs index dbf54b30a..f44b5a8f9 100644 --- a/crates/dbx-cli/src/commands.rs +++ b/crates/dbx-cli/src/commands.rs @@ -103,7 +103,7 @@ fn flag_spec(command: Option<&str>, subcommand: Option<&str>, flag: &str) -> Opt match (command, subcommand, flag) { (Some("conn"), Some("show"), "--redacted") => Some(boolean_flag), (Some("schema"), Some("snapshot"), "--conn" | "--db") => Some(normal_value), - (Some("safe-query"), _, "--conn" | "--db") => Some(normal_value), + (Some("safe-query"), _, "--conn" | "--db" | "--limit") => Some(normal_value), (Some("safe-query"), _, "--sql") => Some(free_text_value), (Some("handoff"), _, "--conn" | "--sql-file") => Some(normal_value), (Some("handoff"), _, "--title" | "--sql" | "--description") => Some(free_text_value), @@ -223,6 +223,29 @@ async fn find_connection( } } +async fn state_with_connection( + config: dbx_core::models::connection::ConnectionConfig, +) -> Result { + let state = open_state().await?; + state.configs.lock().await.insert(config.id.clone(), config.clone()); + + match config.db_type { + dbx_core::models::connection::DatabaseType::Sqlite => { + let pool = dbx_core::db::sqlite::connect_path(&config.host).await?; + state.connections.lock().await.insert(config.id.clone(), dbx_core::connection::PoolKind::Sqlite(pool)); + } + dbx_core::models::connection::DatabaseType::DuckDb => { + let pool = dbx_core::db::duckdb_driver::connect_path(&config.host)?; + state.connections.lock().await.insert(config.id.clone(), dbx_core::connection::PoolKind::DuckDb(pool)); + } + _ => { + state.get_or_create_pool(&config.id, config.database.as_deref()).await?; + } + } + + Ok(state) +} + async fn context(args: &[String]) -> CliEnvelope { if let Err(err) = reject_unexpected_positionals(args, "context", None) { return err; @@ -279,16 +302,22 @@ async fn schema_snapshot(args: &[String]) -> CliEnvelope { return err; } - let Some(_conn) = option_value(args, "--conn") else { + let Some(conn) = option_value(args, "--conn") else { return fail(CliSource::Headless, CliErrorCode::ConnectionNotFound, "--conn is required", true); }; + let config = match find_connection(conn).await { + Ok(config) => config, + Err(err) => return err, + }; + let state = match state_with_connection(config.clone()).await { + Ok(state) => state, + Err(err) => return fail(CliSource::Headless, CliErrorCode::InternalError, err, false), + }; - fail( - CliSource::Headless, - CliErrorCode::InternalError, - "Schema snapshot headless execution is not implemented", - true, - ) + match dbx_core::schema_snapshot::snapshot(&state, &config.id, option_value(args, "--db"), None).await { + Ok(snapshot) => ok(CliSource::Headless, serde_json::to_value(snapshot).unwrap()), + Err(err) => fail(CliSource::Headless, CliErrorCode::InternalError, err, false), + } } async fn safe_query(args: &[String]) -> CliEnvelope { @@ -296,14 +325,64 @@ async fn safe_query(args: &[String]) -> CliEnvelope { return err; } - let Some(_conn) = option_value(args, "--conn") else { + let Some(conn) = option_value(args, "--conn") else { return fail(CliSource::Headless, CliErrorCode::ConnectionNotFound, "--conn is required", true); }; - let Some(_sql) = option_value(args, "--sql") else { + let Some(sql) = option_value(args, "--sql") else { return fail(CliSource::Headless, CliErrorCode::QueryClassificationFailed, "--sql is required", true); }; + let limit = match parse_result_limit(args) { + Ok(limit) => limit as usize, + Err(err) => return err, + }; + 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()); + match risk.operation_class { + dbx_core::sql_safety::OperationClass::Read => {} + dbx_core::sql_safety::OperationClass::Write if risk.is_production => { + return blocked_query(CliErrorCode::ProductionWriteBlocked, &risk); + } + dbx_core::sql_safety::OperationClass::Write => { + return blocked_query(CliErrorCode::HandoffRequired, &risk); + } + dbx_core::sql_safety::OperationClass::Ddl => { + return blocked_query(CliErrorCode::DdlBlocked, &risk); + } + dbx_core::sql_safety::OperationClass::Unknown => { + return blocked_query(CliErrorCode::QueryClassificationFailed, &risk); + } + } + let state = match state_with_connection(config.clone()).await { + Ok(state) => state, + Err(err) => return fail(CliSource::Headless, CliErrorCode::InternalError, err, false), + }; + let database = option_value(args, "--db").or(config.database.as_deref()).unwrap_or(""); - fail(CliSource::Headless, CliErrorCode::InternalError, "Safe query headless execution is not implemented", true) + match dbx_core::query::execute_sql_statement(&state, &config.id, database, sql, None, None).await { + Ok(result) => ok( + CliSource::Headless, + serde_json::json!({ + "risk": risk, + "result": limit_query_result(result, limit), + }), + ), + Err(err) => fail(CliSource::Headless, CliErrorCode::InternalError, err, false), + } +} + +fn blocked_query(code: CliErrorCode, risk: &dbx_core::sql_safety::RiskMetadata) -> CliEnvelope { + fail(CliSource::Headless, code, serde_json::to_string(risk).unwrap(), true) +} + +fn limit_query_result(mut result: dbx_core::db::QueryResult, limit: usize) -> dbx_core::db::QueryResult { + if result.rows.len() > limit { + result.rows.truncate(limit); + result.truncated = true; + } + result } async fn handoff(args: &[String]) -> CliEnvelope { @@ -492,6 +571,30 @@ mod tests { storage.save_connections(configs).await.unwrap(); } + async fn create_sqlite_fixture(path: &std::path::Path) { + 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 teams (id INTEGER PRIMARY KEY, name TEXT NOT NULL)") + .await + .unwrap(); + dbx_core::db::sqlite::execute_query( + &pool, + "CREATE TABLE users (id INTEGER PRIMARY KEY, team_id INTEGER NOT NULL, email TEXT NOT NULL, FOREIGN KEY(team_id) REFERENCES teams(id))", + ) + .await + .unwrap(); + dbx_core::db::sqlite::execute_query(&pool, "INSERT INTO teams (id, name) VALUES (1, 'Core'), (2, 'Data')") + .await + .unwrap(); + dbx_core::db::sqlite::execute_query( + &pool, + "INSERT INTO users (id, team_id, email) VALUES (1, 1, 'ada@example.com'), (2, 2, 'grace@example.com')", + ) + .await + .unwrap(); + pool.close().await; + } + fn success_data(env: CliEnvelope) -> serde_json::Value { match env { CliEnvelope::Success { source, data, .. } => { @@ -782,6 +885,131 @@ mod tests { std::env::remove_var("DBX_APP_DATA_DIR"); } + #[tokio::test] + async fn schema_snapshot_executes_headless_sqlite_snapshot_from_storage() { + let _guard = ENV_LOCK.lock().unwrap(); + let dir = tempfile::tempdir().unwrap(); + let data_path = dir.path().join("fixture.sqlite"); + create_sqlite_fixture(&data_path).await; + std::env::set_var("DBX_APP_DATA_DIR", dir.path()); + seed_connections( + dir.path(), + &[embedded_fixture("sqlite-id", "local-sqlite", DatabaseType::Sqlite, &data_path)], + ) + .await; + + let data = success_data( + dispatch(vec![ + "schema".into(), + "snapshot".into(), + "--conn".into(), + "local-sqlite".into(), + "--format".into(), + "json".into(), + ]) + .await, + ); + + assert_eq!(data["connectionId"], "sqlite-id"); + assert_eq!(data["database"], "main"); + let tables = data["tables"].as_array().expect("tables should be an array"); + assert!(tables.iter().any(|table| table["name"] == "users")); + + std::env::remove_var("DBX_APP_DATA_DIR"); + } + + #[tokio::test] + async fn safe_query_executes_read_sqlite_query_with_limit_and_risk() { + let _guard = ENV_LOCK.lock().unwrap(); + let dir = tempfile::tempdir().unwrap(); + let data_path = dir.path().join("fixture.sqlite"); + create_sqlite_fixture(&data_path).await; + std::env::set_var("DBX_APP_DATA_DIR", dir.path()); + seed_connections( + dir.path(), + &[embedded_fixture("sqlite-id", "local-sqlite", DatabaseType::Sqlite, &data_path)], + ) + .await; + + let data = success_data( + dispatch(vec![ + "safe-query".into(), + "--conn".into(), + "sqlite-id".into(), + "--sql".into(), + "SELECT email FROM users ORDER BY id".into(), + "--limit".into(), + "1".into(), + "--format".into(), + "json".into(), + ]) + .await, + ); + + assert_eq!(data["risk"]["operationClass"], "read"); + assert_eq!(data["result"]["columns"], serde_json::json!(["email"])); + assert_eq!(data["result"]["rows"], serde_json::json!([["ada@example.com"]])); + assert_eq!(data["result"]["truncated"], true); + + std::env::remove_var("DBX_APP_DATA_DIR"); + } + + #[tokio::test] + async fn safe_query_blocks_non_read_sql_with_structured_risk_errors() { + let _guard = ENV_LOCK.lock().unwrap(); + let dir = tempfile::tempdir().unwrap(); + let data_path = dir.path().join("fixture.sqlite"); + create_sqlite_fixture(&data_path).await; + std::env::set_var("DBX_APP_DATA_DIR", dir.path()); + let mut prod = embedded_fixture("prod-sqlite", "prod-sqlite", DatabaseType::Sqlite, &data_path); + prod.color = Some("#ef4444".to_string()); + let dev = embedded_fixture("dev-sqlite", "dev-sqlite", DatabaseType::Sqlite, &data_path); + seed_connections(dir.path(), &[prod, dev]).await; + + let cases = [ + ( + "prod-sqlite", + "UPDATE users SET email = 'x@example.com' WHERE id = 1", + CliErrorCode::ProductionWriteBlocked, + "write", + true, + ), + ( + "dev-sqlite", + "UPDATE users SET email = 'x@example.com' WHERE id = 1", + CliErrorCode::HandoffRequired, + "write", + false, + ), + ("prod-sqlite", "DROP TABLE users", CliErrorCode::DdlBlocked, "ddl", true), + ("prod-sqlite", "VACUUM", CliErrorCode::QueryClassificationFailed, "unknown", true), + ]; + + for (conn, sql, expected_code, expected_class, expected_production) in cases { + match dispatch(vec![ + "safe-query".into(), + "--conn".into(), + conn.into(), + "--sql".into(), + sql.into(), + "--format".into(), + "json".into(), + ]) + .await + { + CliEnvelope::Failure { error, .. } => { + assert_eq!(error.code, expected_code); + let risk: serde_json::Value = serde_json::from_str(&error.message).unwrap(); + assert_eq!(risk["operationClass"], expected_class); + assert_eq!(risk["isProduction"], expected_production); + } + CliEnvelope::Success { .. } => panic!("expected safe-query to block {sql}"), + } + } + + std::env::remove_var("DBX_APP_DATA_DIR"); + } + #[tokio::test] async fn unknown_command_returns_internal_error_envelope() { assert_failure_code( From 2f23df9675809f02d96b2554d77959410e3aa09d Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 19:10:29 +0800 Subject: [PATCH 23/40] =?UTF-8?q?fix(cli):=20=E4=BF=AE=E5=A4=8D=20headless?= =?UTF-8?q?=20=E6=9F=A5=E8=AF=A2=E5=AE=89=E5=85=A8=E4=B8=8E=20limit=20?= =?UTF-8?q?=E4=B8=8B=E6=8E=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 对 schema snapshot 和 safe-query 的内部失败使用安全错误映射 - 为 SQLite/DuckDB 嵌入式连接展开 tilde 路径 - 新增 limit-aware 查询执行接口并让 safe-query 传入读取上限 - 补充 SQLite 大结果集、错误清洗和 tilde 路径回归测试 --- crates/dbx-cli/src/commands.rs | 181 +++++++++++++++++++++++++++---- crates/dbx-core/src/cli.rs | 67 +++++++++++- crates/dbx-core/src/db/sqlite.rs | 15 ++- crates/dbx-core/src/query.rs | 172 +++++++++++++++++++++++++---- 4 files changed, 390 insertions(+), 45 deletions(-) diff --git a/crates/dbx-cli/src/commands.rs b/crates/dbx-cli/src/commands.rs index f44b5a8f9..9756891f7 100644 --- a/crates/dbx-cli/src/commands.rs +++ b/crates/dbx-cli/src/commands.rs @@ -1,4 +1,4 @@ -use dbx_core::cli::{fail, ok, CliEnvelope, CliErrorCode, CliSource}; +use dbx_core::cli::{fail, fail_safe, ok, CliEnvelope, CliErrorCode, CliSource}; const DEFAULT_RESULT_LIMIT: u32 = 50; const MAX_RESULT_LIMIT: u32 = 1000; @@ -188,12 +188,13 @@ fn redact_uri_credentials(value: &str) -> String { async fn load_headless_connections( ) -> Result, CliEnvelope> { - let state = open_state().await.map_err(|err| fail(CliSource::Headless, CliErrorCode::InternalError, err, false))?; + let state = + open_state().await.map_err(|err| fail_safe(CliSource::Headless, CliErrorCode::InternalError, err, false))?; state .storage .load_connections() .await - .map_err(|err| fail(CliSource::Headless, CliErrorCode::InternalError, err, false)) + .map_err(|err| fail_safe(CliSource::Headless, CliErrorCode::InternalError, err, false)) } async fn find_connection( @@ -231,11 +232,13 @@ async fn state_with_connection( match config.db_type { dbx_core::models::connection::DatabaseType::Sqlite => { - let pool = dbx_core::db::sqlite::connect_path(&config.host).await?; + let path = dbx_core::connection::expand_tilde(&config.host); + let pool = dbx_core::db::sqlite::connect_path(&path).await?; state.connections.lock().await.insert(config.id.clone(), dbx_core::connection::PoolKind::Sqlite(pool)); } dbx_core::models::connection::DatabaseType::DuckDb => { - let pool = dbx_core::db::duckdb_driver::connect_path(&config.host)?; + let path = dbx_core::connection::expand_tilde(&config.host); + let pool = dbx_core::db::duckdb_driver::connect_path(&path)?; state.connections.lock().await.insert(config.id.clone(), dbx_core::connection::PoolKind::DuckDb(pool)); } _ => { @@ -311,12 +314,12 @@ async fn schema_snapshot(args: &[String]) -> CliEnvelope { }; let state = match state_with_connection(config.clone()).await { Ok(state) => state, - Err(err) => return fail(CliSource::Headless, CliErrorCode::InternalError, err, false), + Err(err) => return fail_safe(CliSource::Headless, CliErrorCode::InternalError, err, false), }; match dbx_core::schema_snapshot::snapshot(&state, &config.id, option_value(args, "--db"), None).await { Ok(snapshot) => ok(CliSource::Headless, serde_json::to_value(snapshot).unwrap()), - Err(err) => fail(CliSource::Headless, CliErrorCode::InternalError, err, false), + Err(err) => fail_safe(CliSource::Headless, CliErrorCode::InternalError, err, false), } } @@ -357,19 +360,21 @@ async fn safe_query(args: &[String]) -> CliEnvelope { } let state = match state_with_connection(config.clone()).await { Ok(state) => state, - Err(err) => return fail(CliSource::Headless, CliErrorCode::InternalError, err, false), + Err(err) => return fail_safe(CliSource::Headless, CliErrorCode::InternalError, err, false), }; let database = option_value(args, "--db").or(config.database.as_deref()).unwrap_or(""); - match dbx_core::query::execute_sql_statement(&state, &config.id, database, sql, None, None).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!({ "risk": risk, - "result": limit_query_result(result, limit), + "result": result, }), ), - Err(err) => fail(CliSource::Headless, CliErrorCode::InternalError, err, false), + Err(err) => fail_safe(CliSource::Headless, CliErrorCode::InternalError, err, false), } } @@ -377,14 +382,6 @@ fn blocked_query(code: CliErrorCode, risk: &dbx_core::sql_safety::RiskMetadata) fail(CliSource::Headless, code, serde_json::to_string(risk).unwrap(), true) } -fn limit_query_result(mut result: dbx_core::db::QueryResult, limit: usize) -> dbx_core::db::QueryResult { - if result.rows.len() > limit { - result.rows.truncate(limit); - result.truncated = true; - } - result -} - async fn handoff(args: &[String]) -> CliEnvelope { if let Err(err) = reject_unexpected_positionals(args, "handoff", None) { return err; @@ -431,7 +428,7 @@ async fn handoff(args: &[String]) -> CliEnvelope { match queue_handoff(&item).await { Ok(()) => ok(CliSource::Headless, serde_json::json!({ "id": item.id, "status": "queued" })), - Err(err) => fail(CliSource::Headless, CliErrorCode::InternalError, err, false), + Err(err) => fail_safe(CliSource::Headless, CliErrorCode::InternalError, err, false), } } @@ -595,6 +592,26 @@ mod tests { pool.close().await; } + 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(); + 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(); + } + pool.close().await; + } + fn success_data(env: CliEnvelope) -> serde_json::Value { match env { CliEnvelope::Success { source, data, .. } => { @@ -954,6 +971,130 @@ mod tests { std::env::remove_var("DBX_APP_DATA_DIR"); } + #[tokio::test] + async fn safe_query_expands_tilde_for_sqlite_connection_path() { + let _guard = ENV_LOCK.lock().unwrap(); + let dir = tempfile::tempdir().unwrap(); + let home = dir.path().join("home"); + std::fs::create_dir_all(&home).unwrap(); + let data_path = home.join("fixture.sqlite"); + create_sqlite_fixture(&data_path).await; + let previous_home = std::env::var_os("HOME"); + std::env::set_var("HOME", &home); + std::env::set_var("DBX_APP_DATA_DIR", dir.path()); + let mut config = embedded_fixture("sqlite-id", "local-sqlite", DatabaseType::Sqlite, &data_path); + config.host = "~/fixture.sqlite".to_string(); + seed_connections(dir.path(), &[config]).await; + + let data = success_data( + dispatch(vec![ + "safe-query".into(), + "--conn".into(), + "local-sqlite".into(), + "--sql".into(), + "SELECT COUNT(*) FROM users".into(), + "--format".into(), + "json".into(), + ]) + .await, + ); + + assert_eq!(data["result"]["rows"], serde_json::json!([[2]])); + + std::env::remove_var("DBX_APP_DATA_DIR"); + if let Some(previous_home) = previous_home { + std::env::set_var("HOME", previous_home); + } else { + std::env::remove_var("HOME"); + } + } + + #[tokio::test] + async fn safe_query_large_sqlite_result_set_respects_limit() { + let _guard = ENV_LOCK.lock().unwrap(); + let dir = tempfile::tempdir().unwrap(); + let data_path = dir.path().join("large.sqlite"); + create_large_sqlite_fixture(&data_path, 200).await; + std::env::set_var("DBX_APP_DATA_DIR", dir.path()); + seed_connections( + dir.path(), + &[embedded_fixture("sqlite-id", "local-sqlite", DatabaseType::Sqlite, &data_path)], + ) + .await; + + let data = success_data( + dispatch(vec![ + "safe-query".into(), + "--conn".into(), + "sqlite-id".into(), + "--sql".into(), + "SELECT id, value FROM numbers ORDER BY id".into(), + "--limit".into(), + "7".into(), + "--format".into(), + "json".into(), + ]) + .await, + ); + + let rows = data["result"]["rows"].as_array().expect("rows should be an array"); + assert!(rows.len() <= 7, "safe-query returned {} rows for limit 7", rows.len()); + assert_eq!(rows.len(), 7); + assert_eq!(data["result"]["truncated"], true); + + std::env::remove_var("DBX_APP_DATA_DIR"); + } + + #[tokio::test] + async fn schema_snapshot_and_safe_query_sanitize_internal_error_envelopes() { + let _guard = ENV_LOCK.lock().unwrap(); + let dir = tempfile::tempdir().unwrap(); + let private_path = dir.path().join("private").join("missing.sqlite"); + std::env::set_var("DBX_APP_DATA_DIR", dir.path()); + seed_connections( + dir.path(), + &[embedded_fixture("sqlite-id", "local-sqlite", DatabaseType::Sqlite, &private_path)], + ) + .await; + + let cases = [ + dispatch(vec![ + "schema".into(), + "snapshot".into(), + "--conn".into(), + "sqlite-id".into(), + "--format".into(), + "json".into(), + ]) + .await, + dispatch(vec![ + "safe-query".into(), + "--conn".into(), + "sqlite-id".into(), + "--sql".into(), + "SELECT 1".into(), + "--format".into(), + "json".into(), + ]) + .await, + ]; + + for env in cases { + match env { + CliEnvelope::Failure { error, .. } => { + let message = error.message; + assert!(!message.contains(dir.path().to_str().unwrap())); + assert!(!message.contains("missing.sqlite")); + assert!(!message.contains("Database file does not exist")); + assert!(!message.contains("SQLite connection failed")); + } + CliEnvelope::Success { .. } => panic!("expected sanitized failure envelope"), + } + } + + std::env::remove_var("DBX_APP_DATA_DIR"); + } + #[tokio::test] async fn safe_query_blocks_non_read_sql_with_structured_risk_errors() { let _guard = ENV_LOCK.lock().unwrap(); diff --git a/crates/dbx-core/src/cli.rs b/crates/dbx-core/src/cli.rs index dec743430..6296e925a 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")), } @@ -172,6 +176,31 @@ pub fn fail(source: CliSource, code: CliErrorCode, message: impl Into CliEnvelope::Failure { ok: false, source, error: CliError { code, message: message.into(), recoverable } } } +pub fn fail_safe( + source: CliSource, + fallback_code: CliErrorCode, + message: impl AsRef, + recoverable: bool, +) -> CliEnvelope { + let (code, message) = map_safe_error(fallback_code, message.as_ref()); + fail(source, code, message, recoverable) +} + +pub fn map_safe_error(fallback_code: CliErrorCode, message: &str) -> (CliErrorCode, String) { + let lower = message.to_ascii_lowercase(); + if lower.contains("timed out") || lower.contains("timeout") { + return (CliErrorCode::Timeout, "Operation timed out.".to_string()); + } + if lower.contains("connection config not found") || lower == "connection not found" { + return (CliErrorCode::ConnectionNotFound, "Connection not found.".to_string()); + } + if lower.contains("unsupported database") || lower.contains("unsupported database type") { + return (CliErrorCode::UnsupportedDatabaseType, "Unsupported database type.".to_string()); + } + + (fallback_code, "Operation failed. See DBX logs for details.".to_string()) +} + #[cfg(test)] mod tests { use super::*; @@ -279,4 +308,38 @@ mod tests { assert!(err.to_string().contains("ok=true requires data without error")); } + + #[test] + fn maps_timeout_errors_without_leaking_internal_details() { + let (code, message) = map_safe_error( + CliErrorCode::InternalError, + "Query timed out after 30 seconds while reading /Users/alice/private.sqlite", + ); + + assert_eq!(code, CliErrorCode::Timeout); + assert_eq!(message, "Operation timed out."); + assert!(!message.contains("/Users/alice")); + } + + #[test] + fn sanitizes_paths_uri_credentials_and_driver_details() { + let (code, message) = map_safe_error( + CliErrorCode::InternalError, + "SQLite connection failed: Database file does not exist: /Users/alice/private.sqlite; url=mysql://root:secret@localhost/db", + ); + + assert_eq!(code, CliErrorCode::InternalError); + assert_eq!(message, "Operation failed. See DBX logs for details."); + assert!(!message.contains("/Users/alice")); + assert!(!message.contains("root:secret")); + assert!(!message.contains("SQLite connection failed")); + } + + #[test] + fn maps_connection_not_found_to_public_error_code() { + let (code, message) = map_safe_error(CliErrorCode::InternalError, "Connection config not found"); + + assert_eq!(code, CliErrorCode::ConnectionNotFound); + assert_eq!(message, "Connection not found."); + } } diff --git a/crates/dbx-core/src/db/sqlite.rs b/crates/dbx-core/src/db/sqlite.rs index a5b9ef618..f2d668820 100644 --- a/crates/dbx-core/src/db/sqlite.rs +++ b/crates/dbx-core/src/db/sqlite.rs @@ -162,7 +162,16 @@ pub async fn list_triggers(pool: &SqlitePool, _schema: &str, table: &str) -> Res } pub async fn execute_query(pool: &SqlitePool, sql: &str) -> Result { + execute_query_with_row_limit(pool, sql, crate::query::MAX_ROWS).await +} + +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); if starts_with_executable_sql_keyword(sql, &["SELECT", "PRAGMA", "EXPLAIN", "WITH"]) { let desc = pool.describe(sql).await.map_err(|e| e.to_string())?; @@ -191,14 +200,14 @@ pub async fn execute_query(pool: &SqlitePool, sql: &str) -> Result crate::query::MAX_ROWS { + if result_rows.len() > row_limit { break; } } - let truncated = result_rows.len() > crate::query::MAX_ROWS; + let truncated = result_rows.len() > row_limit; if truncated { - result_rows.truncate(crate::query::MAX_ROWS); + result_rows.truncate(row_limit); } Ok(QueryResult { diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 56f9c9a28..df0807c41 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -12,7 +12,16 @@ pub const MAX_ROWS: usize = 10000; pub const QUERY_CANCELED: &str = "Query canceled"; pub fn duckdb_execute(con: &duckdb::Connection, sql: &str) -> Result { + duckdb_execute_with_row_limit(con, sql, MAX_ROWS) +} + +pub fn duckdb_execute_with_row_limit( + con: &duckdb::Connection, + sql: &str, + row_limit: usize, +) -> Result { let start = std::time::Instant::now(); + let row_limit = row_limit.max(1); if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "DESCRIBE", "EXPLAIN", "WITH", "PRAGMA"]) { let mut stmt = con.prepare(sql).map_err(|e| e.to_string())?; @@ -25,9 +34,6 @@ pub fn duckdb_execute(con: &duckdb::Connection, sql: &str) -> Result= MAX_ROWS { - break; - } let vals: Vec = (0..col_count) .map(|i| { row.get::<_, String>(i) @@ -45,9 +51,15 @@ pub fn duckdb_execute(con: &duckdb::Connection, sql: &str) -> Result row_limit { + break; + } } - let truncated = result_rows.len() >= MAX_ROWS; + let truncated = result_rows.len() > row_limit; + if truncated { + result_rows.truncate(row_limit); + } Ok(db::QueryResult { columns, rows: result_rows, @@ -67,9 +79,14 @@ pub fn duckdb_execute(con: &duckdb::Connection, sql: &str) -> Result db::QueryResult { - if result.rows.len() > MAX_ROWS { - result.rows.truncate(MAX_ROWS); +pub fn truncate_result(result: db::QueryResult) -> db::QueryResult { + truncate_result_with_row_limit(result, MAX_ROWS) +} + +pub fn truncate_result_with_row_limit(mut result: db::QueryResult, row_limit: usize) -> db::QueryResult { + let row_limit = row_limit.max(1); + if result.rows.len() > row_limit { + result.rows.truncate(row_limit); result.truncated = true; } result @@ -142,6 +159,18 @@ pub async fn do_execute( schema: Option<&str>, cancel_token: Option, ) -> Result { + do_execute_with_row_limit(state, pool_key, sql, schema, cancel_token, MAX_ROWS).await +} + +pub async fn do_execute_with_row_limit( + state: &AppState, + pool_key: &str, + sql: &str, + schema: Option<&str>, + cancel_token: Option, + row_limit: usize, +) -> Result { + let row_limit = row_limit.max(1); let connections = state.connections.read().await; let pool = connections.get(pool_key).ok_or("Connection not found")?; @@ -153,7 +182,7 @@ pub async fn do_execute( wait_for_query(cancel_token, async move { let task = tokio::task::spawn_blocking(move || { let con = con.lock().map_err(|e| e.to_string())?; - duckdb_execute(&con, &sql) + duckdb_execute_with_row_limit(&con, &sql, row_limit) }); task.await.map_err(|e| e.to_string())? }) @@ -178,7 +207,7 @@ pub async fn do_execute( PoolKind::Sqlite(p) => { let p = p.clone(); drop(connections); - wait_for_query(cancel_token, db::sqlite::execute_query(&p, sql)).await + wait_for_query(cancel_token, db::sqlite::execute_query_with_row_limit(&p, sql, row_limit)).await } PoolKind::ClickHouse(client) => { let client = client.clone(); @@ -186,7 +215,7 @@ pub async fn do_execute( drop(connections); wait_for_query(cancel_token, db::clickhouse_driver::execute_query(&client, &database, sql)) .await - .map(truncate_result) + .map(|result| truncate_result_with_row_limit(result, row_limit)) } PoolKind::SqlServer(client) => { let client = client.clone(); @@ -199,7 +228,9 @@ pub async fn do_execute( }, None => client.lock().await, }; - wait_for_query(cancel_token, db::sqlserver::execute_query(&mut client, sql)).await.map(truncate_result) + wait_for_query(cancel_token, db::sqlserver::execute_query(&mut client, sql)) + .await + .map(|result| truncate_result_with_row_limit(result, row_limit)) } PoolKind::Oracle(pool) => { let client = pool.client(); @@ -218,9 +249,11 @@ pub async fn do_execute( if let Some(schema) = schema { wait_for_query(cancel_token, db::oracle_driver::execute_query_with_schema(&*client, &schema, sql)) .await - .map(truncate_result) + .map(|result| truncate_result_with_row_limit(result, row_limit)) } else { - wait_for_query(cancel_token, db::oracle_driver::execute_query(&*client, sql)).await.map(truncate_result) + wait_for_query(cancel_token, db::oracle_driver::execute_query(&*client, sql)) + .await + .map(|result| truncate_result_with_row_limit(result, row_limit)) } } PoolKind::Elasticsearch(client) => { @@ -229,7 +262,7 @@ pub async fn do_execute( drop(connections); wait_for_query(cancel_token, db::elasticsearch_driver::execute_rest_query(&client, &sql)) .await - .map(truncate_result) + .map(|result| truncate_result_with_row_limit(result, row_limit)) } PoolKind::Redis(_) => Err("Use Redis-specific commands".to_string()), PoolKind::MongoDb(_) => Err("Use MongoDB-specific commands".to_string()), @@ -250,7 +283,7 @@ pub async fn do_execute( task.await.map_err(|e| e.to_string())? }) .await - .map(truncate_result) + .map(|result| truncate_result_with_row_limit(result, row_limit)) } PoolKind::Gaussdb(client) => { let client = client.clone(); @@ -261,7 +294,7 @@ pub async fn do_execute( db::gaussdb_driver::execute_query(&mut client, &sql).await }) .await - .map(truncate_result) + .map(|result| truncate_result_with_row_limit(result, row_limit)) } PoolKind::ExternalTabular(ext_pool) => { if !starts_with_executable_sql_keyword(sql, &["SELECT", "WITH", "SHOW", "DESCRIBE", "EXPLAIN", "PRAGMA"]) { @@ -273,7 +306,7 @@ pub async fn do_execute( wait_for_query(cancel_token, async move { let task = tokio::task::spawn_blocking(move || { let con = con.lock().map_err(|e| e.to_string())?; - duckdb_execute(&con, &sql) + duckdb_execute_with_row_limit(&con, &sql, row_limit) }); task.await.map_err(|e| e.to_string())? }) @@ -294,7 +327,7 @@ pub async fn do_execute( session.invoke::("executeQuery", params).await }) .await - .map(truncate_result) + .map(|result| truncate_result_with_row_limit(result, row_limit)) } } } @@ -307,6 +340,19 @@ pub async fn execute_sql_statement( schema: Option<&str>, cancel_token: Option, ) -> Result { + execute_sql_statement_with_row_limit(state, connection_id, database, sql, schema, cancel_token, MAX_ROWS).await +} + +pub async fn execute_sql_statement_with_row_limit( + state: &AppState, + connection_id: &str, + database: &str, + sql: &str, + schema: Option<&str>, + cancel_token: Option, + row_limit: usize, +) -> Result { + let row_limit = row_limit.max(1); let pool_key = if database.is_empty() { connection_id.to_string() } else { @@ -317,13 +363,13 @@ pub async fn execute_sql_statement( return Err(canceled_error()); } - let result = do_execute(state, &pool_key, sql, schema, cancel_token.clone()).await; + let result = do_execute_with_row_limit(state, &pool_key, sql, schema, cancel_token.clone(), row_limit).await; match &result { Err(e) if is_connection_error(e) && !is_canceled(&cancel_token) => { let db_opt = if database.is_empty() { None } else { Some(database) }; let new_key = state.reconnect_pool(connection_id, db_opt).await?; - do_execute(state, &new_key, sql, schema, cancel_token).await + do_execute_with_row_limit(state, &new_key, sql, schema, cancel_token, row_limit).await } _ => result, } @@ -746,6 +792,9 @@ async fn exec_tx_none_inner( #[cfg(test)] mod tests { use super::*; + use crate::connection::AppState; + use crate::models::connection::{ConnectionConfig, DatabaseType}; + use crate::storage::Storage; #[tokio::test] async fn wait_for_query_returns_cancelled_when_token_is_cancelled() { @@ -816,4 +865,87 @@ mod tests { assert!(!is_connection_error("syntax error at position 5")); assert!(!is_connection_error("os error 13")); } + + fn sqlite_config(path: &std::path::Path) -> ConnectionConfig { + ConnectionConfig { + id: "sqlite-id".to_string(), + name: "local-sqlite".to_string(), + db_type: DatabaseType::Sqlite, + driver_profile: None, + driver_label: None, + url_params: None, + host: path.display().to_string(), + port: 0, + username: String::new(), + password: String::new(), + database: None, + color: None, + ssh_enabled: false, + ssh_host: String::new(), + ssh_port: 22, + ssh_user: String::new(), + ssh_password: String::new(), + ssh_key_path: String::new(), + ssh_key_passphrase: String::new(), + ssh_expose_lan: false, + ssh_connect_timeout_secs: crate::models::connection::default_ssh_connect_timeout_secs(), + proxy_enabled: false, + proxy_type: crate::models::connection::ProxyType::Socks5, + proxy_host: String::new(), + proxy_port: 1080, + proxy_username: String::new(), + proxy_password: String::new(), + ssl: false, + sysdba: false, + connection_string: None, + external_config: None, + jdbc_driver_class: None, + jdbc_driver_paths: Vec::new(), + } + } + + async fn sqlite_state_with_rows(row_count: usize) -> (AppState, std::path::PathBuf, std::path::PathBuf) { + let unique = uuid::Uuid::new_v4(); + let data_path = std::env::temp_dir().join(format!("dbx-query-data-{unique}.sqlite")); + let storage_path = std::env::temp_dir().join(format!("dbx-query-storage-{unique}.sqlite")); + std::fs::File::create(&data_path).unwrap(); + let pool = db::sqlite::connect_path(&data_path.display().to_string()).await.unwrap(); + db::sqlite::execute_query(&pool, "CREATE TABLE numbers (id INTEGER PRIMARY KEY)").await.unwrap(); + for id in 1..=row_count { + sqlx::query("INSERT INTO numbers (id) VALUES (?)").bind(id as i64).execute(&pool).await.unwrap(); + } + pool.close().await; + + let storage = Storage::open(&storage_path).await.unwrap(); + let state = AppState::new(storage); + let config = sqlite_config(&data_path); + state.configs.write().await.insert(config.id.clone(), config); + let pool = db::sqlite::connect_path(&data_path.display().to_string()).await.unwrap(); + state.connections.write().await.insert("sqlite-id".to_string(), PoolKind::Sqlite(pool)); + (state, data_path, storage_path) + } + + #[tokio::test] + async fn execute_sql_statement_with_row_limit_caps_sqlite_rows() { + let (state, data_path, storage_path) = sqlite_state_with_rows(200).await; + + let result = execute_sql_statement_with_row_limit( + &state, + "sqlite-id", + "", + "SELECT id FROM numbers ORDER BY id", + None, + None, + 7, + ) + .await + .unwrap(); + + assert!(result.rows.len() <= 7); + assert_eq!(result.rows.len(), 7); + assert!(result.truncated); + + let _ = std::fs::remove_file(data_path); + let _ = std::fs::remove_file(storage_path); + } } From 2c3704a365ef4a814f3db18dbf780507a6115a6c Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 19:52:32 +0800 Subject: [PATCH 24/40] =?UTF-8?q?fix(query):=20=E7=BB=9F=E4=B8=80=E6=89=A7?= =?UTF-8?q?=E8=A1=8C=E5=85=A5=E5=8F=A3=E7=BB=93=E6=9E=9C=E8=A1=8C=E6=95=B0?= =?UTF-8?q?=E6=88=AA=E6=96=AD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 在 execute_sql_statement_with_row_limit 成功返回前统一应用 row_limit - 补充 truncate_result_with_row_limit 的截断、保留和零值限制单测 --- crates/dbx-core/src/query.rs | 42 ++++++++++++++++++++++++++++++++++-- 1 file changed, 40 insertions(+), 2 deletions(-) diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index df0807c41..41663707e 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -365,14 +365,16 @@ pub async fn execute_sql_statement_with_row_limit( let result = do_execute_with_row_limit(state, &pool_key, sql, schema, cancel_token.clone(), row_limit).await; - match &result { + let result = match &result { Err(e) if is_connection_error(e) && !is_canceled(&cancel_token) => { let db_opt = if database.is_empty() { None } else { Some(database) }; let new_key = state.reconnect_pool(connection_id, db_opt).await?; do_execute_with_row_limit(state, &new_key, sql, schema, cancel_token, row_limit).await } _ => result, - } + }; + + result.map(|result| truncate_result_with_row_limit(result, row_limit)) } pub async fn execute_multi_core( @@ -796,6 +798,42 @@ mod tests { use crate::models::connection::{ConnectionConfig, DatabaseType}; use crate::storage::Storage; + fn query_result_with_rows(row_count: usize, truncated: bool) -> db::QueryResult { + db::QueryResult { + columns: vec!["id".to_string()], + rows: (1..=row_count).map(|id| vec![serde_json::Value::Number(id.into())]).collect(), + affected_rows: 0, + execution_time_ms: 0, + truncated, + } + } + + #[test] + fn truncate_result_with_row_limit_truncates_rows_and_marks_result() { + let result = truncate_result_with_row_limit(query_result_with_rows(5, false), 3); + + assert_eq!(result.rows.len(), 3); + assert!(result.truncated); + assert_eq!(result.rows[0][0], serde_json::Value::Number(1.into())); + assert_eq!(result.rows[2][0], serde_json::Value::Number(3.into())); + } + + #[test] + fn truncate_result_with_row_limit_keeps_untruncated_result_when_within_limit() { + let result = truncate_result_with_row_limit(query_result_with_rows(3, false), 3); + + assert_eq!(result.rows.len(), 3); + assert!(!result.truncated); + } + + #[test] + fn truncate_result_with_row_limit_treats_zero_limit_as_one() { + let result = truncate_result_with_row_limit(query_result_with_rows(2, false), 0); + + assert_eq!(result.rows.len(), 1); + assert!(result.truncated); + } + #[tokio::test] async fn wait_for_query_returns_cancelled_when_token_is_cancelled() { let token = CancellationToken::new(); From 4267c6f1e02d2786b3c99613cdac028a26b1f248 Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 19:53:04 +0800 Subject: [PATCH 25/40] =?UTF-8?q?style:=20=E5=BA=94=E7=94=A8=20CLI=20?= =?UTF-8?q?=E7=9B=B8=E5=85=B3=20Rust=20=E6=A0=BC=E5=BC=8F=E5=8C=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- crates/dbx-cli/src/commands.rs | 17 +++++--------- crates/dbx-cli/src/main.rs | 5 +++- crates/dbx-cli/src/runtime_client.rs | 34 +++++++++++++++++++++------- crates/dbx-core/src/cli.rs | 8 ++----- crates/dbx-core/src/db/sqlite.rs | 6 +---- 5 files changed, 39 insertions(+), 31 deletions(-) diff --git a/crates/dbx-cli/src/commands.rs b/crates/dbx-cli/src/commands.rs index 9756891f7..349ff0e2d 100644 --- a/crates/dbx-cli/src/commands.rs +++ b/crates/dbx-cli/src/commands.rs @@ -364,9 +364,7 @@ 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!({ @@ -595,19 +593,16 @@ 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; } diff --git a/crates/dbx-cli/src/main.rs b/crates/dbx-cli/src/main.rs index 9a5a19356..f34308a75 100644 --- a/crates/dbx-cli/src/main.rs +++ b/crates/dbx-cli/src/main.rs @@ -4,7 +4,10 @@ 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 effa954a0..81dda04c8 100644 --- a/crates/dbx-cli/src/runtime_client.rs +++ b/crates/dbx-cli/src/runtime_client.rs @@ -17,7 +17,8 @@ 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") @@ -42,30 +43,44 @@ 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 { @@ -85,7 +100,10 @@ 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 6296e925a..534c4e1f4 100644 --- a/crates/dbx-core/src/cli.rs +++ b/crates/dbx-core/src/cli.rs @@ -116,12 +116,8 @@ 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 f2d668820..a82651323 100644 --- a/crates/dbx-core/src/db/sqlite.rs +++ b/crates/dbx-core/src/db/sqlite.rs @@ -165,11 +165,7 @@ 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); From 8a81bdfdd56bb4fa1688f694aad9590eb25082b8 Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 20:00:52 +0800 Subject: [PATCH 26/40] =?UTF-8?q?feat(core):=20=E6=94=AF=E6=8C=81=20MySQL?= =?UTF-8?q?=20=E5=92=8C=20Postgres=20=E6=9F=A5=E8=AF=A2=E8=A1=8C=E6=95=B0?= =?UTF-8?q?=E9=99=90=E5=88=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 新增 MySQL/Postgres limit-aware 查询执行函数 - 保持旧执行 API 默认转调 MAX_ROWS - 让 query.rs 对 MySQL/Postgres 传递 row_limit - 补充 limit-aware 驱动调用回归测试 --- crates/dbx-core/src/db/mysql.rs | 27 ++++++++++++++----- crates/dbx-core/src/db/postgres.rs | 42 +++++++++++++++++++++++++----- crates/dbx-core/src/query.rs | 19 +++++++++++--- 3 files changed, 73 insertions(+), 15 deletions(-) diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index 5946f0436..3fa858591 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -249,7 +249,17 @@ pub async fn get_columns(pool: &MySqlPool, database: &str, table: &str) -> Resul } pub async fn execute_query(pool: &MySqlPool, sql: &str, bare: bool) -> Result { + execute_query_with_row_limit(pool, sql, bare, crate::query::MAX_ROWS).await +} + +pub async fn execute_query_with_row_limit( + pool: &MySqlPool, + sql: &str, + bare: bool, + row_limit: usize, +) -> Result { let start = Instant::now(); + let row_limit = row_limit.max(1); if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "DESCRIBE", "EXPLAIN"]) { if bare { @@ -269,14 +279,14 @@ pub async fn execute_query(pool: &MySqlPool, sql: &str, bare: bool) -> Result crate::query::MAX_ROWS { + if result_rows.len() > row_limit { break; } } - let truncated = result_rows.len() > crate::query::MAX_ROWS; + let truncated = result_rows.len() > row_limit; if truncated { - result_rows.truncate(crate::query::MAX_ROWS); + result_rows.truncate(row_limit); } Ok(QueryResult { @@ -301,14 +311,14 @@ pub async fn execute_query(pool: &MySqlPool, sql: &str, bare: bool) -> Result crate::query::MAX_ROWS { + if result_rows.len() > row_limit { break; } } - let truncated = result_rows.len() > crate::query::MAX_ROWS; + let truncated = result_rows.len() > row_limit; if truncated { - result_rows.truncate(crate::query::MAX_ROWS); + result_rows.truncate(row_limit); } Ok(QueryResult { @@ -413,6 +423,11 @@ pub async fn list_triggers(pool: &MySqlPool, database: &str, table: &str) -> Res mod tests { use super::*; + #[allow(dead_code)] + fn mysql_limit_aware_execute_query_api_compiles(pool: &MySqlPool, sql: &str, bare: bool) { + let _ = execute_query_with_row_limit(pool, sql, bare, 7); + } + #[test] fn numeric_metadata_accepts_unsigned_information_schema_values() { assert_eq!(numeric_metadata_u64_to_i32(Some(65)), Some(65)); diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index bec71c4fd..c7891fe16 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -281,7 +281,12 @@ pub async fn get_columns(pool: &PgPool, schema: &str, table: &str) -> Result Result { + execute_query_with_row_limit(pool, sql, crate::query::MAX_ROWS).await +} + +pub async fn execute_query_with_row_limit(pool: &PgPool, sql: &str, row_limit: usize) -> Result { let start = Instant::now(); + let row_limit = row_limit.max(1); if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "EXPLAIN", "WITH", "TABLE"]) { let mut stream = sqlx::query(sql).persistent(false).fetch(pool); @@ -301,7 +306,7 @@ pub async fn execute_query(pool: &PgPool, sql: &str) -> Result crate::query::MAX_ROWS { + if result_rows.len() > row_limit { break; } } @@ -311,9 +316,9 @@ pub async fn execute_query(pool: &PgPool, sql: &str) -> Result crate::query::MAX_ROWS; + let truncated = result_rows.len() > row_limit; if truncated { - result_rows.truncate(crate::query::MAX_ROWS); + result_rows.truncate(row_limit); } Ok(QueryResult { @@ -337,11 +342,21 @@ pub async fn execute_query(pool: &PgPool, sql: &str) -> Result Result { + execute_query_with_schema_and_row_limit(pool, schema, sql, crate::query::MAX_ROWS).await +} + +pub async fn execute_query_with_schema_and_row_limit( + pool: &PgPool, + schema: &str, + sql: &str, + row_limit: usize, +) -> Result { let mut conn = pool.acquire().await.map_err(|e| e.to_string())?; let set_path = format!("SET search_path TO \"{}\", public", schema); sqlx::query(&set_path).execute(&mut *conn).await.map_err(|e| e.to_string())?; let start = Instant::now(); + let row_limit = row_limit.max(1); if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "EXPLAIN", "WITH", "TABLE"]) { let mut stream = sqlx::query(sql).persistent(false).fetch(&mut *conn); @@ -361,7 +376,7 @@ pub async fn execute_query_with_schema(pool: &PgPool, schema: &str, sql: &str) - .map(|i| pg_value_to_json(&row, i, column_types.get(i).map(String::as_str).unwrap_or(""))) .collect(), ); - if result_rows.len() > crate::query::MAX_ROWS { + if result_rows.len() > row_limit { break; } } @@ -372,9 +387,9 @@ pub async fn execute_query_with_schema(pool: &PgPool, schema: &str, sql: &str) - columns = desc.columns().iter().map(|c| c.name().to_string()).collect(); } - let truncated = result_rows.len() > crate::query::MAX_ROWS; + let truncated = result_rows.len() > row_limit; if truncated { - result_rows.truncate(crate::query::MAX_ROWS); + result_rows.truncate(row_limit); } Ok(QueryResult { @@ -446,6 +461,21 @@ pub async fn list_indexes(pool: &PgPool, schema: &str, table: &str) -> Result Result, String> { let rows: Vec = sqlx::query( "SELECT kcu.constraint_name, kcu.column_name, \ diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 41663707e..d007e9403 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -192,16 +192,20 @@ pub async fn do_execute_with_row_limit( let p = p.clone(); let bare = *mode == crate::connection::MysqlMode::Bare; drop(connections); - wait_for_query(cancel_token, db::mysql::execute_query(&p, sql, bare)).await + wait_for_query(cancel_token, db::mysql::execute_query_with_row_limit(&p, sql, bare, row_limit)).await } PoolKind::Postgres(p) => { let p = p.clone(); let schema = schema.map(|s| s.to_string()); drop(connections); if let Some(schema) = schema { - wait_for_query(cancel_token, db::postgres::execute_query_with_schema(&p, &schema, sql)).await + wait_for_query( + cancel_token, + db::postgres::execute_query_with_schema_and_row_limit(&p, &schema, sql, row_limit), + ) + .await } else { - wait_for_query(cancel_token, db::postgres::execute_query(&p, sql)).await + wait_for_query(cancel_token, db::postgres::execute_query_with_row_limit(&p, sql, row_limit)).await } } PoolKind::Sqlite(p) => { @@ -834,6 +838,15 @@ mod tests { assert!(result.truncated); } + #[test] + fn mysql_and_postgres_pool_branches_call_limit_aware_drivers() { + let source = include_str!("query.rs"); + + assert!(source.contains("db::mysql::execute_query_with_row_limit(&p, sql, bare, row_limit)")); + assert!(source.contains("db::postgres::execute_query_with_schema_and_row_limit(&p, &schema, sql, row_limit)")); + assert!(source.contains("db::postgres::execute_query_with_row_limit(&p, sql, row_limit)")); + } + #[tokio::test] async fn wait_for_query_returns_cancelled_when_token_is_cancelled() { let token = CancellationToken::new(); From 9425658ecd4b5d690e4212959c3557369d182771 Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 20:13:03 +0800 Subject: [PATCH 27/40] =?UTF-8?q?feat(agent-runtime):=20=E6=96=B0=E5=A2=9E?= =?UTF-8?q?=20Tauri=20agent=20runtime=20=E6=9C=8D=E5=8A=A1?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 启动独立 localhost server 并写入 owner-only discovery - 提供 token 认证的 context、selection、result 和 handoff 端点 - 接入 Tauri 启动、invoke 命令和退出清理 - 补充认证、路由、discovery 权限与 symlink 防护测试 --- src-tauri/src/commands/agent_runtime.rs | 374 ++++++++++++++++++++++++ src-tauri/src/commands/mod.rs | 1 + src-tauri/src/lib.rs | 10 + 3 files changed, 385 insertions(+) create mode 100644 src-tauri/src/commands/agent_runtime.rs diff --git a/src-tauri/src/commands/agent_runtime.rs b/src-tauri/src/commands/agent_runtime.rs new file mode 100644 index 000000000..a1d1e7de1 --- /dev/null +++ b/src-tauri/src/commands/agent_runtime.rs @@ -0,0 +1,374 @@ +use serde::{Deserialize, Serialize}; +use std::fs::OpenOptions; +use std::io::Write; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use tauri::{AppHandle, Manager}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; +use tokio::sync::{oneshot, RwLock}; + +use super::connection::AppState; + +const BIND_ADDR: &str = "127.0.0.1:0"; +const DISCOVERY_FILE: &str = "agent-runtime.json"; + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AgentRuntimeSnapshot { + pub active_connection_id: Option, + pub active_connection_name: Option, + pub database: Option, + pub schema: Option, + pub active_tab_id: Option, + pub active_tab_title: Option, + pub sql: Option, + pub selected_sql: Option, + pub selection: Option, + pub result: Option, +} + +#[derive(Clone)] +pub struct AgentRuntimeState { + pub token: String, + pub snapshot: Arc>, + pub handoffs: Arc>>, +} + +pub struct AgentRuntimeServer { + state: AgentRuntimeState, + discovery_path: PathBuf, + shutdown: std::sync::Mutex>>, +} + +impl AgentRuntimeServer { + pub fn state(&self) -> &AgentRuntimeState { + &self.state + } + + pub fn cleanup(&self) { + if let Ok(mut shutdown) = self.shutdown.lock() { + if let Some(tx) = shutdown.take() { + let _ = tx.send(()); + } + } + cleanup_discovery_file(&self.discovery_path); + } +} + +impl Drop for AgentRuntimeServer { + fn drop(&mut self) { + if let Ok(mut shutdown) = self.shutdown.lock() { + if let Some(tx) = shutdown.take() { + let _ = tx.send(()); + } + } + cleanup_discovery_file(&self.discovery_path); + } +} + +#[derive(Debug, PartialEq, Eq)] +struct RuntimeResponse { + status: &'static str, + body: serde_json::Value, +} + +#[tauri::command] +pub async fn agent_runtime_update_snapshot( + runtime: tauri::State<'_, AgentRuntimeServer>, + snapshot: AgentRuntimeSnapshot, +) -> Result<(), String> { + *runtime.state().snapshot.write().await = snapshot; + Ok(()) +} + +#[tauri::command] +pub async fn agent_runtime_load_handoffs( + app_state: tauri::State<'_, Arc>, + 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()); + Ok(items) +} + +pub fn start(app: AppHandle) -> AgentRuntimeServer { + let token = uuid::Uuid::new_v4().to_string(); + let state = AgentRuntimeState { + token: token.clone(), + snapshot: Arc::new(RwLock::new(AgentRuntimeSnapshot::default())), + handoffs: Arc::new(RwLock::new(Vec::new())), + }; + let discovery_path = + app.path().app_data_dir().map(|dir| dir.join(DISCOVERY_FILE)).unwrap_or_else(|_| PathBuf::from(DISCOVERY_FILE)); + let (shutdown_tx, shutdown_rx) = oneshot::channel(); + let server_state = state.clone(); + + tauri::async_runtime::spawn(async move { + run_server(app, server_state, shutdown_rx).await; + }); + + AgentRuntimeServer { state, discovery_path, shutdown: std::sync::Mutex::new(Some(shutdown_tx)) } +} + +async fn run_server(app: AppHandle, state: AgentRuntimeState, mut shutdown: oneshot::Receiver<()>) { + let listener = match TcpListener::bind(BIND_ADDR).await { + Ok(listener) => listener, + Err(err) => { + log::warn!("Agent runtime failed to bind {BIND_ADDR}: {err}"); + return; + } + }; + let port = listener.local_addr().map(|addr| addr.port()).unwrap_or(0); + let discovery_path = match app.path().app_data_dir() { + Ok(dir) => match write_discovery_file(&dir, port, &state.token) { + Ok(path) => Some(path), + Err(err) => { + log::warn!("Agent runtime discovery write failed: {err}"); + None + } + }, + Err(err) => { + log::warn!("Agent runtime app data dir unavailable: {err}"); + None + } + }; + log::info!("Agent runtime listening on 127.0.0.1:{port}"); + + loop { + tokio::select! { + _ = &mut shutdown => { + if let Some(path) = discovery_path.as_deref() { + cleanup_discovery_file(path); + } + break; + } + accepted = listener.accept() => { + let Ok((stream, _)) = accepted else { continue }; + let st = state.clone(); + tokio::spawn(async move { + handle_connection(stream, st).await; + }); + } + } + } +} + +async fn handle_connection(mut stream: TcpStream, state: AgentRuntimeState) { + let mut buf = vec![0u8; 65536]; + let Ok(n) = stream.read(&mut buf).await else { + return; + }; + if n == 0 { + return; + } + + let request = String::from_utf8_lossy(&buf[..n]); + if !is_authorized(&request, &state.token) { + respond_json(&mut stream, "401 Unauthorized", serde_json::json!({"error": "unauthorized"})).await; + return; + } + + let first_line = request.lines().next().unwrap_or(""); + let body = request.split("\r\n\r\n").nth(1).unwrap_or(""); + let response = route_request(first_line, body, &state).await; + respond_json(&mut stream, response.status, response.body).await; +} + +fn is_authorized(request: &str, token: &str) -> bool { + request.lines().any(|line| line.trim_end() == format!("Authorization: Bearer {token}")) +} + +async fn route_request(first_line: &str, body: &str, state: &AgentRuntimeState) -> RuntimeResponse { + if first_line.starts_with("GET /context ") || first_line.starts_with("GET /context?") { + return RuntimeResponse { + status: "200 OK", + body: serde_json::to_value(&*state.snapshot.read().await).unwrap_or_else(|_| serde_json::json!({})), + }; + } + + if first_line.starts_with("GET /selection ") || first_line.starts_with("GET /selection?") { + let snapshot = state.snapshot.read().await; + return RuntimeResponse { + status: "200 OK", + body: snapshot.selection.clone().unwrap_or_else(|| serde_json::json!({"type": "none"})), + }; + } + + if first_line.starts_with("GET /result/current ") || first_line.starts_with("GET /result/current?") { + let snapshot = state.snapshot.read().await; + return RuntimeResponse { + status: "200 OK", + body: snapshot.result.clone().unwrap_or_else(|| serde_json::json!({"columns": [], "rows": []})), + }; + } + + if first_line.starts_with("POST /handoff ") { + let mut item = match serde_json::from_str::(body) { + Ok(item) => item, + Err(_) => { + return RuntimeResponse { + status: "400 Bad Request", + body: serde_json::json!({"error": "invalid handoff"}), + }; + } + }; + item.status = dbx_core::handoff::HandoffStatus::Shown; + let id = item.id.clone(); + state.handoffs.write().await.push(item); + return RuntimeResponse { status: "200 OK", body: serde_json::json!({"id": id, "status": "shown"}) }; + } + + RuntimeResponse { status: "404 Not Found", body: serde_json::json!({"error": "not found"}) } +} + +fn write_discovery_file(dir: &Path, port: u16, token: &str) -> Result { + std::fs::create_dir_all(dir).map_err(|err| err.to_string())?; + let path = dir.join(DISCOVERY_FILE); + if let Ok(metadata) = std::fs::symlink_metadata(&path) { + if metadata.file_type().is_symlink() || !metadata.is_file() { + std::fs::remove_file(&path).map_err(|err| err.to_string())?; + } + } + let payload = serde_json::json!({ "port": port, "token": token }); + let body = serde_json::to_vec(&payload).map_err(|err| err.to_string())?; + + let mut options = OpenOptions::new(); + options.create(true).write(true).truncate(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + options.mode(0o600); + } + let mut file = options.open(&path).map_err(|err| err.to_string())?; + file.write_all(&body).map_err(|err| err.to_string())?; + file.sync_all().map_err(|err| err.to_string())?; + + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + let mut permissions = file.metadata().map_err(|err| err.to_string())?.permissions(); + permissions.set_mode(0o600); + std::fs::set_permissions(&path, permissions).map_err(|err| err.to_string())?; + } + + Ok(path) +} + +fn cleanup_discovery_file(path: &Path) { + if path.exists() { + let _ = std::fs::remove_file(path); + } +} + +async fn respond_json(stream: &mut TcpStream, status: &str, body: serde_json::Value) { + let body = serde_json::to_string(&body).unwrap_or_else(|_| "{}".to_string()); + let resp = format!( + "HTTP/1.1 {status}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ); + let _ = stream.write_all(resp.as_bytes()).await; +} + +#[cfg(test)] +mod tests { + use super::*; + + fn runtime_state() -> AgentRuntimeState { + AgentRuntimeState { + token: "secret-token".to_string(), + snapshot: Arc::new(RwLock::new(AgentRuntimeSnapshot::default())), + handoffs: Arc::new(RwLock::new(Vec::new())), + } + } + + #[cfg(unix)] + fn mode(path: &Path) -> u32 { + use std::os::unix::fs::PermissionsExt; + + std::fs::metadata(path).unwrap().permissions().mode() & 0o777 + } + + #[test] + fn authorization_requires_exact_bearer_token() { + assert!(is_authorized("GET /context HTTP/1.1\r\nAuthorization: Bearer secret-token\r\n\r\n", "secret-token",)); + assert!(!is_authorized("GET /context HTTP/1.1\r\nAuthorization: Bearer wrong\r\n\r\n", "secret-token",)); + assert!(!is_authorized("GET /context HTTP/1.1\r\n\r\n", "secret-token")); + } + + #[tokio::test] + async fn routes_context_selection_result_and_handoff_from_shared_state() { + let state = runtime_state(); + *state.snapshot.write().await = AgentRuntimeSnapshot { + active_connection_id: Some("conn-1".to_string()), + active_connection_name: Some("Local".to_string()), + selection: Some(serde_json::json!({"type": "grid-cells", "cells": [[1]]})), + result: Some(serde_json::json!({"columns": ["id"], "rows": [[1]]})), + ..AgentRuntimeSnapshot::default() + }; + + let context = route_request("GET /context HTTP/1.1", "", &state).await; + assert_eq!(context.status, "200 OK"); + assert_eq!(context.body["activeConnectionId"], "conn-1"); + + let selection = route_request("GET /selection HTTP/1.1", "", &state).await; + assert_eq!(selection.status, "200 OK"); + assert_eq!(selection.body["type"], "grid-cells"); + + let result = route_request("GET /result/current?limit=50 HTTP/1.1", "", &state).await; + assert_eq!(result.status, "200 OK"); + assert_eq!(result.body["columns"][0], "id"); + + 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 handoff = route_request("POST /handoff HTTP/1.1", &serde_json::to_string(&item).unwrap(), &state).await; + assert_eq!(handoff.status, "200 OK"); + assert_eq!(handoff.body["id"], item.id); + assert_eq!(state.handoffs.read().await.len(), 1); + } + + #[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())); + std::fs::create_dir_all(&dir).unwrap(); + + let path = write_discovery_file(&dir, 4321, "secret-token").unwrap(); + let payload: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap(); + assert_eq!(payload["port"], 4321); + assert_eq!(payload["token"], "secret-token"); + #[cfg(unix)] + assert_eq!(mode(&path), 0o600); + + cleanup_discovery_file(&path); + assert!(!path.exists()); + + let _ = std::fs::remove_dir_all(dir); + } + + #[cfg(unix)] + #[test] + fn discovery_file_replaces_existing_symlink() { + let dir = std::env::temp_dir().join(format!("dbx-agent-runtime-symlink-test-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let target = dir.join("target.json"); + let link = dir.join(DISCOVERY_FILE); + std::fs::write(&target, "{}").unwrap(); + std::os::unix::fs::symlink(&target, &link).unwrap(); + + let path = write_discovery_file(&dir, 4321, "secret-token").unwrap(); + + assert!(!std::fs::symlink_metadata(&path).unwrap().file_type().is_symlink()); + assert_eq!(mode(&path), 0o600); + + let _ = std::fs::remove_dir_all(dir); + } +} diff --git a/src-tauri/src/commands/mod.rs b/src-tauri/src/commands/mod.rs index 57442c8b2..69f61492d 100644 --- a/src-tauri/src/commands/mod.rs +++ b/src-tauri/src/commands/mod.rs @@ -1,3 +1,4 @@ +pub mod agent_runtime; pub mod ai; pub mod connection; #[allow(dead_code, unused_imports)] diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index cc5c97a11..7ce3f2cda 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -39,6 +39,8 @@ pub fn run() { let app_handle = app.handle().clone(); commands::mcp_bridge::start(app_handle, state); + let runtime_state = commands::agent_runtime::start(app.handle().clone()); + app.manage(runtime_state); #[cfg(not(target_os = "macos"))] { @@ -57,6 +59,8 @@ pub fn run() { } }) .invoke_handler(tauri::generate_handler![ + commands::agent_runtime::agent_runtime_update_snapshot, + commands::agent_runtime::agent_runtime_load_handoffs, commands::ai::ai_complete, commands::ai::ai_stream, commands::ai::ai_cancel_stream, @@ -145,6 +149,12 @@ pub fn run() { .build(tauri::generate_context!()) .expect("error while building tauri application") .run(|app_handle, event| { + if let RunEvent::ExitRequested { .. } = &event { + if let Some(runtime) = app_handle.try_state::() { + runtime.cleanup(); + } + } + #[cfg(target_os = "macos")] if let RunEvent::Reopen { has_visible_windows, .. } = &event { if !has_visible_windows { From 67a7ef4746e892c896a08eb21fde1aabaa134c7a Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 20:19:42 +0800 Subject: [PATCH 28/40] =?UTF-8?q?fix(agent-runtime):=20=E4=BF=AE=E5=A4=8D?= =?UTF-8?q?=E8=BF=90=E8=A1=8C=E6=97=B6=20HTTP=20=E8=AF=BB=E5=8F=96?= =?UTF-8?q?=E4=B8=8E=E5=8F=91=E7=8E=B0=E6=96=87=E4=BB=B6=E5=86=99=E5=85=A5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- Cargo.lock | 1 + src-tauri/Cargo.toml | 1 + src-tauri/src/commands/agent_runtime.rs | 272 +++++++++++++++++++++--- 3 files changed, 250 insertions(+), 24 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index daf4f1613..ea6031c52 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1644,6 +1644,7 @@ dependencies = [ "dbx-core", "duckdb", "futures", + "libc", "log", "mongodb", "percent-encoding", diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index 08b00aa33..d4ec4bf0e 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -48,4 +48,5 @@ russh = "0.60" csv = "1.4.0" calamine = "0.30.1" zip = "4.6.1" +libc = "0.2" dbx-core = { path = "../crates/dbx-core" } diff --git a/src-tauri/src/commands/agent_runtime.rs b/src-tauri/src/commands/agent_runtime.rs index a1d1e7de1..a96ec9178 100644 --- a/src-tauri/src/commands/agent_runtime.rs +++ b/src-tauri/src/commands/agent_runtime.rs @@ -12,6 +12,8 @@ use super::connection::AppState; const BIND_ADDR: &str = "127.0.0.1:0"; const DISCOVERY_FILE: &str = "agent-runtime.json"; +const MAX_HEADER_BYTES: usize = 16 * 1024; +const MAX_BODY_BYTES: usize = 1024 * 1024; #[derive(Debug, Clone, Default, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] @@ -73,6 +75,13 @@ struct RuntimeResponse { body: serde_json::Value, } +#[derive(Debug)] +struct RuntimeRequest { + first_line: String, + headers: Vec<(String, String)>, + body: String, +} + #[tauri::command] pub async fn agent_runtime_update_snapshot( runtime: tauri::State<'_, AgentRuntimeServer>, @@ -155,28 +164,124 @@ async fn run_server(app: AppHandle, state: AgentRuntimeState, mut shutdown: ones } async fn handle_connection(mut stream: TcpStream, state: AgentRuntimeState) { - let mut buf = vec![0u8; 65536]; - let Ok(n) = stream.read(&mut buf).await else { - return; + let request = match read_request(&mut stream).await { + Ok(Some(request)) => request, + Ok(None) => return, + Err(response) => { + respond_json(&mut stream, response.status, response.body).await; + return; + } }; - if n == 0 { - return; - } - let request = String::from_utf8_lossy(&buf[..n]); - if !is_authorized(&request, &state.token) { + if !is_authorized_headers(&request.headers, &state.token) { respond_json(&mut stream, "401 Unauthorized", serde_json::json!({"error": "unauthorized"})).await; return; } - let first_line = request.lines().next().unwrap_or(""); - let body = request.split("\r\n\r\n").nth(1).unwrap_or(""); - let response = route_request(first_line, body, &state).await; + let response = route_request(&request.first_line, &request.body, &state).await; respond_json(&mut stream, response.status, response.body).await; } +async fn read_request(stream: &mut TcpStream) -> Result, RuntimeResponse> { + let mut buf = Vec::new(); + let header_end = loop { + if let Some(pos) = find_header_end(&buf) { + break pos; + } + if buf.len() >= MAX_HEADER_BYTES { + return Err(RuntimeResponse { + status: "431 Request Header Fields Too Large", + body: serde_json::json!({"error": "headers too large"}), + }); + } + + let mut chunk = [0u8; 8192]; + let n = stream.read(&mut chunk).await.map_err(|_| RuntimeResponse { + status: "400 Bad Request", + body: serde_json::json!({"error": "invalid request"}), + })?; + if n == 0 { + return if buf.is_empty() { + Ok(None) + } else { + Err(RuntimeResponse { + status: "400 Bad Request", + body: serde_json::json!({"error": "incomplete request"}), + }) + }; + } + buf.extend_from_slice(&chunk[..n]); + }; + + let header_text = String::from_utf8_lossy(&buf[..header_end]); + let mut lines = header_text.lines(); + let first_line = lines.next().unwrap_or("").to_string(); + let headers: Vec<(String, String)> = lines + .filter_map(|line| { + let (name, value) = line.split_once(':')?; + Some((name.trim().to_string(), value.trim().to_string())) + }) + .collect(); + let content_length = content_length(&headers)?; + if content_length > MAX_BODY_BYTES { + return Err(RuntimeResponse { + status: "413 Payload Too Large", + body: serde_json::json!({"error": "body too large"}), + }); + } + + let body_start = header_end + 4; + let body_end = body_start + content_length; + while buf.len() < body_end { + let mut chunk = [0u8; 8192]; + let n = stream.read(&mut chunk).await.map_err(|_| RuntimeResponse { + status: "400 Bad Request", + body: serde_json::json!({"error": "invalid request"}), + })?; + if n == 0 { + return Err(RuntimeResponse { + status: "400 Bad Request", + body: serde_json::json!({"error": "incomplete body"}), + }); + } + buf.extend_from_slice(&chunk[..n]); + } + + Ok(Some(RuntimeRequest { + first_line, + headers, + body: String::from_utf8_lossy(&buf[body_start..body_end]).to_string(), + })) +} + +fn find_header_end(buf: &[u8]) -> Option { + buf.windows(4).position(|window| window == b"\r\n\r\n") +} + +fn content_length(headers: &[(String, String)]) -> Result { + match headers.iter().find(|(name, _)| name.eq_ignore_ascii_case("content-length")) { + Some((_, value)) => value.parse::().map_err(|_| RuntimeResponse { + status: "400 Bad Request", + body: serde_json::json!({"error": "invalid content-length"}), + }), + None => Ok(0), + } +} + +#[cfg(test)] fn is_authorized(request: &str, token: &str) -> bool { - request.lines().any(|line| line.trim_end() == format!("Authorization: Bearer {token}")) + request.lines().any(|line| { + let Some((name, value)) = line.split_once(':') else { + return false; + }; + name.trim().eq_ignore_ascii_case("authorization") && value.trim() == format!("Bearer {token}") + }) +} + +fn is_authorized_headers(headers: &[(String, String)], token: &str) -> bool { + headers + .iter() + .any(|(name, value)| name.eq_ignore_ascii_case("authorization") && value == &format!("Bearer {token}")) } async fn route_request(first_line: &str, body: &str, state: &AgentRuntimeState) -> RuntimeResponse { @@ -197,10 +302,11 @@ async fn route_request(first_line: &str, body: &str, state: &AgentRuntimeState) if first_line.starts_with("GET /result/current ") || first_line.starts_with("GET /result/current?") { let snapshot = state.snapshot.read().await; - return RuntimeResponse { - status: "200 OK", - body: snapshot.result.clone().unwrap_or_else(|| serde_json::json!({"columns": [], "rows": []})), - }; + let mut body = snapshot.result.clone().unwrap_or_else(|| serde_json::json!({"columns": [], "rows": []})); + if let Some(limit) = query_limit(first_line) { + truncate_result_rows(&mut body, limit); + } + return RuntimeResponse { status: "200 OK", body }; } if first_line.starts_with("POST /handoff ") { @@ -222,25 +328,37 @@ async fn route_request(first_line: &str, body: &str, state: &AgentRuntimeState) RuntimeResponse { status: "404 Not Found", body: serde_json::json!({"error": "not found"}) } } +fn query_limit(first_line: &str) -> Option { + let target = first_line.split_whitespace().nth(1)?; + let query = target.split_once('?')?.1; + query.split('&').find_map(|pair| { + let (key, value) = pair.split_once('=')?; + (key == "limit").then(|| value.parse::().ok()).flatten() + }) +} + +fn truncate_result_rows(result: &mut serde_json::Value, limit: usize) { + if let Some(rows) = result.get_mut("rows").and_then(|rows| rows.as_array_mut()) { + rows.truncate(limit); + } +} + fn write_discovery_file(dir: &Path, port: u16, token: &str) -> Result { std::fs::create_dir_all(dir).map_err(|err| err.to_string())?; let path = dir.join(DISCOVERY_FILE); - if let Ok(metadata) = std::fs::symlink_metadata(&path) { - if metadata.file_type().is_symlink() || !metadata.is_file() { - std::fs::remove_file(&path).map_err(|err| err.to_string())?; - } - } + let temp_path = dir.join(format!("{DISCOVERY_FILE}.{}.tmp", uuid::Uuid::new_v4())); let payload = serde_json::json!({ "port": port, "token": token }); let body = serde_json::to_vec(&payload).map_err(|err| err.to_string())?; let mut options = OpenOptions::new(); - options.create(true).write(true).truncate(true); + options.create_new(true).write(true); #[cfg(unix)] { use std::os::unix::fs::OpenOptionsExt; options.mode(0o600); + options.custom_flags(libc::O_NOFOLLOW); } - let mut file = options.open(&path).map_err(|err| err.to_string())?; + let mut file = options.open(&temp_path).map_err(|err| err.to_string())?; file.write_all(&body).map_err(|err| err.to_string())?; file.sync_all().map_err(|err| err.to_string())?; @@ -249,7 +367,13 @@ fn write_discovery_file(dir: &Path, port: u16, token: &str) -> Result AgentRuntimeState { AgentRuntimeState { @@ -289,13 +414,45 @@ mod tests { std::fs::metadata(path).unwrap().permissions().mode() & 0o777 } + async fn serve_once(state: AgentRuntimeState) -> std::net::SocketAddr { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + handle_connection(stream, state).await; + }); + addr + } + #[test] fn authorization_requires_exact_bearer_token() { assert!(is_authorized("GET /context HTTP/1.1\r\nAuthorization: Bearer secret-token\r\n\r\n", "secret-token",)); + assert!(is_authorized("GET /context HTTP/1.1\r\nauthorization: Bearer secret-token\r\n\r\n", "secret-token",)); assert!(!is_authorized("GET /context HTTP/1.1\r\nAuthorization: Bearer wrong\r\n\r\n", "secret-token",)); assert!(!is_authorized("GET /context HTTP/1.1\r\n\r\n", "secret-token")); } + #[tokio::test] + async fn accepts_reqwest_lowercase_authorization_header() { + let state = runtime_state(); + *state.snapshot.write().await = AgentRuntimeSnapshot { + active_connection_id: Some("conn-1".to_string()), + ..AgentRuntimeSnapshot::default() + }; + let addr = serve_once(state).await; + + let response = reqwest::Client::new() + .get(format!("http://{addr}/context")) + .bearer_auth("secret-token") + .send() + .await + .unwrap(); + + assert_eq!(response.status(), reqwest::StatusCode::OK); + let body: serde_json::Value = response.json().await.unwrap(); + assert_eq!(body["activeConnectionId"], "conn-1"); + } + #[tokio::test] async fn routes_context_selection_result_and_handoff_from_shared_state() { let state = runtime_state(); @@ -336,6 +493,73 @@ mod tests { assert_eq!(state.handoffs.read().await.len(), 1); } + #[tokio::test] + async fn result_current_limit_truncates_rows() { + let state = runtime_state(); + *state.snapshot.write().await = AgentRuntimeSnapshot { + result: Some(serde_json::json!({"columns": ["id"], "rows": [[1], [2], [3]]})), + ..AgentRuntimeSnapshot::default() + }; + + let result = route_request("GET /result/current?limit=2 HTTP/1.1", "", &state).await; + + assert_eq!(result.status, "200 OK"); + assert_eq!(result.body["rows"], serde_json::json!([[1], [2]])); + } + + #[tokio::test] + async fn reads_fragmented_handoff_body_until_content_length() { + let state = runtime_state(); + let addr = serve_once(state.clone()).await; + let item = dbx_core::handoff::HandoffItem::queued( + "conn-1".to_string(), + "Local".to_string(), + Some("main".to_string()), + "Review SQL".to_string(), + None, + "select ".to_string() + &"1".repeat(70_000), + dbx_core::sql_safety::OperationClass::Read, + dbx_core::sql_safety::RiskLevel::Low, + false, + ); + let body = serde_json::to_string(&item).unwrap(); + let head = format!( + "POST /handoff HTTP/1.1\r\nHost: {addr}\r\nAuthorization: Bearer secret-token\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n", + body.len() + ); + let split_at = body.len() / 2; + let mut stream = TcpStream::connect(addr).await.unwrap(); + + stream.write_all(head.as_bytes()).await.unwrap(); + stream.write_all(body[..split_at].as_bytes()).await.unwrap(); + tokio::task::yield_now().await; + stream.write_all(body[split_at..].as_bytes()).await.unwrap(); + let mut response = Vec::new(); + stream.read_to_end(&mut response).await.unwrap(); + + let response = String::from_utf8(response).unwrap(); + assert!(response.starts_with("HTTP/1.1 200 OK"), "{response}"); + assert_eq!(state.handoffs.read().await.len(), 1); + } + + #[tokio::test] + async fn rejects_body_larger_than_limit() { + let state = runtime_state(); + let addr = serve_once(state).await; + let body_len = 1_048_577; + let request = format!( + "POST /handoff HTTP/1.1\r\nHost: {addr}\r\nAuthorization: Bearer secret-token\r\nContent-Length: {body_len}\r\n\r\n" + ); + let mut stream = TcpStream::connect(addr).await.unwrap(); + + stream.write_all(request.as_bytes()).await.unwrap(); + let mut response = Vec::new(); + stream.read_to_end(&mut response).await.unwrap(); + + let response = String::from_utf8(response).unwrap(); + assert!(response.starts_with("HTTP/1.1 413 Payload Too Large"), "{response}"); + } + #[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())); From 5b62f936d1a15619f4d0c73ba2fbc0375853e0cd Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 20:20:26 +0800 Subject: [PATCH 29/40] =?UTF-8?q?style:=20=E5=BA=94=E7=94=A8=20Rust=20?= =?UTF-8?q?=E6=A0=BC=E5=BC=8F=E5=8C=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- crates/dbx-cli/src/commands.rs | 17 +++++++++----- crates/dbx-cli/src/main.rs | 5 +--- crates/dbx-cli/src/runtime_client.rs | 34 +++++++--------------------- crates/dbx-core/src/cli.rs | 8 +++++-- crates/dbx-core/src/db/sqlite.rs | 6 ++++- crates/dbx-core/src/storage.rs | 15 ++++++------ 6 files changed, 38 insertions(+), 47 deletions(-) diff --git a/crates/dbx-cli/src/commands.rs b/crates/dbx-cli/src/commands.rs index 349ff0e2d..9756891f7 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!({ @@ -593,16 +595,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; } diff --git a/crates/dbx-cli/src/main.rs b/crates/dbx-cli/src/main.rs index f34308a75..9a5a19356 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 81dda04c8..effa954a0 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 534c4e1f4..6296e925a 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 a82651323..f2d668820 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 347caec01..444792949 100644 --- a/crates/dbx-core/src/storage.rs +++ b/crates/dbx-core/src/storage.rs @@ -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(()); @@ -1057,10 +1058,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")); } From feff43c459204193e2cd7d55033cd2ac0c468c91 Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 21:20:59 +0800 Subject: [PATCH 30/40] =?UTF-8?q?style:=20=E6=81=A2=E5=A4=8D=20hook=20?= =?UTF-8?q?=E5=90=8E=20Rust=20=E6=A0=BC=E5=BC=8F=E5=8C=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- crates/dbx-cli/src/commands.rs | 17 +++++-------- crates/dbx-cli/src/main.rs | 5 +++- crates/dbx-cli/src/runtime_client.rs | 34 +++++++++++++++++++------ crates/dbx-core/src/cli.rs | 8 ++---- crates/dbx-core/src/db/sqlite.rs | 6 +---- crates/dbx-core/src/storage.rs | 15 ++++++----- src-tauri/src/commands/agent_runtime.rs | 5 +++- src-tauri/src/commands/mod.rs | 2 +- 8 files changed, 52 insertions(+), 40 deletions(-) diff --git a/crates/dbx-cli/src/commands.rs b/crates/dbx-cli/src/commands.rs index 9756891f7..349ff0e2d 100644 --- a/crates/dbx-cli/src/commands.rs +++ b/crates/dbx-cli/src/commands.rs @@ -364,9 +364,7 @@ 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!({ @@ -595,19 +593,16 @@ 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; } diff --git a/crates/dbx-cli/src/main.rs b/crates/dbx-cli/src/main.rs index 9a5a19356..f34308a75 100644 --- a/crates/dbx-cli/src/main.rs +++ b/crates/dbx-cli/src/main.rs @@ -4,7 +4,10 @@ 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 effa954a0..81dda04c8 100644 --- a/crates/dbx-cli/src/runtime_client.rs +++ b/crates/dbx-cli/src/runtime_client.rs @@ -17,7 +17,8 @@ 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") @@ -42,30 +43,44 @@ 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 { @@ -85,7 +100,10 @@ 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 6296e925a..534c4e1f4 100644 --- a/crates/dbx-core/src/cli.rs +++ b/crates/dbx-core/src/cli.rs @@ -116,12 +116,8 @@ 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 f2d668820..a82651323 100644 --- a/crates/dbx-core/src/db/sqlite.rs +++ b/crates/dbx-core/src/db/sqlite.rs @@ -165,11 +165,7 @@ 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 444792949..347caec01 100644 --- a/crates/dbx-core/src/storage.rs +++ b/crates/dbx-core/src/storage.rs @@ -144,11 +144,10 @@ 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(()); @@ -1058,8 +1057,10 @@ 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")); } diff --git a/src-tauri/src/commands/agent_runtime.rs b/src-tauri/src/commands/agent_runtime.rs index a96ec9178..8c1d093b2 100644 --- a/src-tauri/src/commands/agent_runtime.rs +++ b/src-tauri/src/commands/agent_runtime.rs @@ -306,7 +306,10 @@ 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 ") { diff --git a/src-tauri/src/commands/mod.rs b/src-tauri/src/commands/mod.rs index 69f61492d..2040e4ca0 100644 --- a/src-tauri/src/commands/mod.rs +++ b/src-tauri/src/commands/mod.rs @@ -1,5 +1,5 @@ -pub mod agent_runtime; pub mod ai; +pub mod agent_runtime; pub mod connection; #[allow(dead_code, unused_imports)] mod connection_secrets; From c144330d7913b6b5f2b0f6bacc32c70099a1e33d Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 21:21:22 +0800 Subject: [PATCH 31/40] =?UTF-8?q?feat(runtime):=20=E5=90=8C=E6=AD=A5?= =?UTF-8?q?=E5=89=8D=E7=AB=AF=E7=8A=B6=E6=80=81=E5=88=B0=20agent=20runtime?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 新增 snapshot 构造 helper 和 agent runtime store - 同步 active tab、连接、库、schema、SQL、选中 SQL、表格选择和结果采样 - 补充 snapshot payload 组装测试 --- src/App.vue | 10 +++- src/components/grid/DataGrid.vue | 30 ++++++++++- src/lib/agentRuntimeSnapshot.ts | 62 +++++++++++++++++++++++ src/lib/api.ts | 4 ++ src/lib/http.ts | 7 +++ src/lib/tauri.ts | 9 ++++ src/stores/agentRuntimeStore.ts | 51 +++++++++++++++++++ src/stores/queryStore.ts | 21 ++++++++ tests/agentRuntimeSnapshot.test.ts | 80 ++++++++++++++++++++++++++++++ 9 files changed, 271 insertions(+), 3 deletions(-) create mode 100644 src/lib/agentRuntimeSnapshot.ts create mode 100644 src/stores/agentRuntimeStore.ts create mode 100644 tests/agentRuntimeSnapshot.test.ts diff --git a/src/App.vue b/src/App.vue index bf8558af8..d849be045 100644 --- a/src/App.vue +++ b/src/App.vue @@ -15,6 +15,7 @@ import UpdateDialog from "@/components/layout/UpdateDialog.vue"; import LoginPage from "@/components/auth/LoginPage.vue"; import { useConnectionStore } from "@/stores/connectionStore"; import { useQueryStore } from "@/stores/queryStore"; +import { useAgentRuntimeStore } from "@/stores/agentRuntimeStore"; import { useSettingsStore } from "@/stores/settingsStore"; import { useSavedSqlStore } from "@/stores/savedSqlStore"; import { useToast } from "@/composables/useToast"; @@ -46,6 +47,7 @@ import type { HistoryEntry } from "@/lib/tauri"; const { t } = useI18n(); const connectionStore = useConnectionStore(); const queryStore = useQueryStore(); +const agentRuntimeStore = useAgentRuntimeStore(); const settingsStore = useSettingsStore(); const savedSqlStore = useSavedSqlStore(); const { message: toastMessage, visible: toastVisible, toast } = useToast(); @@ -158,6 +160,7 @@ watch( () => queryStore.activeTabId, () => { selectedSql.value = ""; + agentRuntimeStore.setSelectedSql(""); activeOutputView.value = "result"; }, ); @@ -614,7 +617,12 @@ onUnmounted(() => { if (queryStore.activeTabId) queryStore.updateSql(queryStore.activeTabId, v); } " - @editor-selection-change="(v: string) => (selectedSql = v)" + @editor-selection-change=" + (v: string) => { + selectedSql = v; + agentRuntimeStore.setSelectedSql(v); + } + " @editor-cursor-change="(p: number) => (cursorPos = p)" @format-error="toast(t('toolbar.formatSqlFailed'))" @reload=" diff --git a/src/components/grid/DataGrid.vue b/src/components/grid/DataGrid.vue index 9e5a6e644..3b115df73 100644 --- a/src/components/grid/DataGrid.vue +++ b/src/components/grid/DataGrid.vue @@ -62,9 +62,17 @@ import type { QueryResult, ColumnInfo, DatabaseType } from "@/types/database"; import * as api from "@/lib/api"; import { buildTableSelectSql, quoteTableIdentifier } from "@/lib/tableSelectSql"; import { isHiddenGridColumn, usesSyntheticRowIdKey } from "@/lib/tableEditing"; -import { formatGridSqlLiteral } from "@/lib/dataGridSql"; -import { matchesRowStatusFilter, type RowStatus, type RowStatusFilter } from "@/lib/gridRowStatus"; import { displayCellValue, type CellValue } from "@/lib/cellValue"; +import { buildDataGridSaveStatements, formatGridSqlLiteral } from "@/lib/dataGridSql"; +import { formatMarkdownTable } from "@/lib/markdownTable"; +import { buildXlsxWorkbook } from "@/lib/xlsxExport"; +import { + matchesRowStatusFilter, + rowStatusFilterAfterAddingRow, + type RowStatus, + type RowStatusFilter, +} from "@/lib/gridRowStatus"; +import { useAgentRuntimeStore } from "@/stores/agentRuntimeStore"; import { useToast } from "@/composables/useToast"; import { useDataGridExport } from "@/composables/useDataGridExport"; @@ -74,6 +82,7 @@ import { useDataGridEditor } from "@/composables/useDataGridEditor"; const { t } = useI18n(); const { toast } = useToast(); +const agentRuntimeStore = useAgentRuntimeStore(); const props = defineProps<{ result: QueryResult; @@ -1090,6 +1099,22 @@ const activeCellDetail = computed(() => { const detailEditValue = ref(""); const isEditingDetail = ref(false); +watch( + selectedCells, + (data) => { + agentRuntimeStore.setSelection( + selectedCellCount.value > 0 + ? { + type: "grid-cells", + range: selectedRange.value, + data, + } + : { type: "none" }, + ); + }, + { deep: true }, +); + function startDetailEdit() { const detail = activeCellDetail.value; if (!detail || !detail.isEditable) return; @@ -1611,6 +1636,7 @@ watch( onUnmounted(() => { cleanupFrames(); + agentRuntimeStore.setSelection({ type: "none" }); onDdlResizeEnd(); finishCellSelection(); clearTimeout(_searchTimer); diff --git a/src/lib/agentRuntimeSnapshot.ts b/src/lib/agentRuntimeSnapshot.ts new file mode 100644 index 000000000..5494863fd --- /dev/null +++ b/src/lib/agentRuntimeSnapshot.ts @@ -0,0 +1,62 @@ +import type { ConnectionConfig, QueryResult, QueryTab } from "@/types/database"; + +export const DEFAULT_RESULT_SAMPLE_LIMIT = 50; + +export interface RuntimeResultSample { + columns: string[]; + rows: QueryResult["rows"]; + truncated: boolean; + executionTimeMs: number; + sampleLimit: number; +} + +export interface AgentRuntimeSnapshot { + activeConnectionId?: string; + activeConnectionName?: string; + database?: string; + schema?: string; + activeTabId?: string; + activeTabTitle?: string; + sql?: string; + selectedSql?: string; + selection: unknown; + result?: RuntimeResultSample; +} + +export interface BuildAgentRuntimeSnapshotOptions { + tabs: QueryTab[]; + activeTabId: string | null; + getConnection: (connectionId: string) => Pick | undefined; + selectedSql?: string; + selection?: unknown; + resultSampleLimit?: number; +} + +export function buildAgentRuntimeSnapshot(options: BuildAgentRuntimeSnapshotOptions): AgentRuntimeSnapshot { + const tab = options.tabs.find((item) => item.id === options.activeTabId); + const conn = tab ? options.getConnection(tab.connectionId) : undefined; + const selectedSql = options.selectedSql?.trim(); + const limit = options.resultSampleLimit ?? DEFAULT_RESULT_SAMPLE_LIMIT; + const result = tab?.result + ? { + columns: tab.result.columns, + rows: tab.result.rows.slice(0, limit), + truncated: tab.result.rows.length > limit || !!tab.result.truncated, + executionTimeMs: tab.result.execution_time_ms, + sampleLimit: limit, + } + : undefined; + + return { + activeConnectionId: tab?.connectionId, + activeConnectionName: conn?.name, + database: tab?.database, + schema: tab?.schema, + activeTabId: tab?.id, + activeTabTitle: tab?.title, + sql: tab?.sql, + selectedSql: selectedSql || undefined, + selection: options.selection ?? { type: "none" }, + result, + }; +} diff --git a/src/lib/api.ts b/src/lib/api.ts index 86fa229dd..0098c528d 100644 --- a/src/lib/api.ts +++ b/src/lib/api.ts @@ -72,6 +72,10 @@ export const executeScript = forward("executeScript"); export const executeInTransaction = forward("executeInTransaction"); export const cancelQuery = forward("cancelQuery"); +// Agent Runtime +export const agentRuntimeUpdateSnapshot = forward("agentRuntimeUpdateSnapshot"); +export const agentRuntimeLoadHandoffs = forward("agentRuntimeLoadHandoffs"); + // AI export const aiComplete = forward("aiComplete"); export const aiStream = forward("aiStream"); diff --git a/src/lib/http.ts b/src/lib/http.ts index 6fdfb977a..148a68f71 100644 --- a/src/lib/http.ts +++ b/src/lib/http.ts @@ -39,6 +39,7 @@ import type { TableImportSummary, TableImportProgress, } from "./tauri"; +import type { AgentRuntimeSnapshot } from "@/lib/agentRuntimeSnapshot"; // --------------------------------------------------------------------------- // Helpers @@ -290,6 +291,12 @@ export async function cancelQuery(executionId: string): Promise { return post("/api/query/cancel", { executionId }); } +export async function agentRuntimeUpdateSnapshot(_snapshot: AgentRuntimeSnapshot): Promise {} + +export async function agentRuntimeLoadHandoffs(): Promise { + return []; +} + // --------------------------------------------------------------------------- // AI // --------------------------------------------------------------------------- diff --git a/src/lib/tauri.ts b/src/lib/tauri.ts index 1c2906ffc..2be0e47a6 100644 --- a/src/lib/tauri.ts +++ b/src/lib/tauri.ts @@ -19,6 +19,7 @@ import type { SavedSqlFolder, SavedSqlLibrary, } from "@/types/database"; +import type { AgentRuntimeSnapshot } from "@/lib/agentRuntimeSnapshot"; import type { AiConfig } from "@/stores/settingsStore"; export interface AiMessage { @@ -80,6 +81,14 @@ export async function loadAiConfig(): Promise { return invoke("load_ai_config"); } +export async function agentRuntimeUpdateSnapshot(snapshot: AgentRuntimeSnapshot): Promise { + return invoke("agent_runtime_update_snapshot", { snapshot }); +} + +export async function agentRuntimeLoadHandoffs(): Promise { + return invoke("agent_runtime_load_handoffs"); +} + // --- AI Conversations --- export interface AiChatMessage { diff --git a/src/stores/agentRuntimeStore.ts b/src/stores/agentRuntimeStore.ts new file mode 100644 index 000000000..97d7f38be --- /dev/null +++ b/src/stores/agentRuntimeStore.ts @@ -0,0 +1,51 @@ +import { defineStore } from "pinia"; +import { ref } from "vue"; +import * as api from "@/lib/api"; +import { buildAgentRuntimeSnapshot, DEFAULT_RESULT_SAMPLE_LIMIT } from "@/lib/agentRuntimeSnapshot"; +import { useConnectionStore } from "@/stores/connectionStore"; +import { useQueryStore } from "@/stores/queryStore"; + +export const useAgentRuntimeStore = defineStore("agentRuntime", () => { + const selection = ref({ type: "none" }); + const selectedSql = ref(""); + let timer: ReturnType | null = null; + + function setSelection(value: unknown) { + selection.value = value; + scheduleSync(); + } + + function setSelectedSql(value: string) { + selectedSql.value = value; + scheduleSync(); + } + + function scheduleSync() { + if (timer) clearTimeout(timer); + timer = setTimeout(() => { + timer = null; + void syncNow(); + }, 100); + } + + async function syncNow() { + const connectionStore = useConnectionStore(); + const queryStore = useQueryStore(); + const snapshot = buildAgentRuntimeSnapshot({ + tabs: queryStore.tabs, + activeTabId: queryStore.activeTabId, + getConnection: (connectionId) => connectionStore.getConfig(connectionId), + selectedSql: selectedSql.value, + selection: selection.value, + resultSampleLimit: DEFAULT_RESULT_SAMPLE_LIMIT, + }); + + try { + await api.agentRuntimeUpdateSnapshot(snapshot); + } catch (err) { + console.debug("[DBX] Agent runtime snapshot sync skipped:", err); + } + } + + return { selection, selectedSql, setSelection, setSelectedSql, scheduleSync, syncNow }; +}); diff --git a/src/stores/queryStore.ts b/src/stores/queryStore.ts index d3094ddb9..4532c5b42 100644 --- a/src/stores/queryStore.ts +++ b/src/stores/queryStore.ts @@ -10,6 +10,7 @@ import { analyzeEditableQuery, allPrimaryKeysPresent } from "@/lib/sqlAnalysis"; import { restoreOpenTabsState, serializeOpenTabs } from "@/lib/openTabsPersistence"; import * as api from "@/lib/api"; import { useConnectionStore } from "@/stores/connectionStore"; +import { useAgentRuntimeStore } from "@/stores/agentRuntimeStore"; import { isTauriRuntime } from "@/lib/tauriRuntime"; import type { SavedSqlFile } from "@/types/database"; @@ -72,11 +73,18 @@ export const useQueryStore = defineStore("query", () => { return tabs.value.find((t) => t.connectionId === connectionId && t.database === database && t.title === title); } + function scheduleAgentRuntimeSync() { + useAgentRuntimeStore().scheduleSync(); + } + + watch(activeTabId, () => scheduleAgentRuntimeSync()); + function createTab(connectionId: string, database: string, title?: string, mode: QueryTab["mode"] = "query") { if (title) { const existing = findTabByTitle(connectionId, database, title); if (existing) { activeTabId.value = existing.id; + scheduleAgentRuntimeSync(); return existing.id; } } @@ -95,6 +103,7 @@ export const useQueryStore = defineStore("query", () => { }; tabs.value.push(tab); activeTabId.value = id; + scheduleAgentRuntimeSync(); return id; } @@ -145,6 +154,7 @@ export const useQueryStore = defineStore("query", () => { if (activeTabId.value === id) { activeTabId.value = tabs.value[Math.min(idx, tabs.value.length - 1)]?.id ?? null; } + scheduleAgentRuntimeSync(); } function closeOtherTabs(id: string) { @@ -153,6 +163,7 @@ export const useQueryStore = defineStore("query", () => { const next = closeOtherTabsState(tabs.value, activeTabId.value, id); tabs.value = next.tabs; activeTabId.value = next.activeTabId; + scheduleAgentRuntimeSync(); } function closeAllTabs() { @@ -161,6 +172,7 @@ export const useQueryStore = defineStore("query", () => { const next = closeAllTabsState(tabs.value, activeTabId.value); tabs.value = next.tabs; activeTabId.value = next.activeTabId; + scheduleAgentRuntimeSync(); } function updateSql(id: string, sql: string) { @@ -169,6 +181,7 @@ export const useQueryStore = defineStore("query", () => { tab.sql = sql; tab.resultSortedSql = undefined; tab.resultBaseSql = undefined; + scheduleAgentRuntimeSync(); } } @@ -183,6 +196,7 @@ export const useQueryStore = defineStore("query", () => { const existing = tabs.value.find((tab) => tab.savedSqlId === file.id); if (existing) { activeTabId.value = existing.id; + scheduleAgentRuntimeSync(); return existing.id; } @@ -202,6 +216,7 @@ export const useQueryStore = defineStore("query", () => { }; tabs.value.push(tab); activeTabId.value = id; + scheduleAgentRuntimeSync(); return id; } @@ -224,6 +239,7 @@ export const useQueryStore = defineStore("query", () => { tab.resultSortedSql = undefined; clearExplain(tab); tab.tableMeta = undefined; + scheduleAgentRuntimeSync(); } function updateSchema(id: string, schema: string | undefined) { @@ -231,6 +247,7 @@ export const useQueryStore = defineStore("query", () => { if (!tab || tab.schema === schema) return; tab.schema = schema; if (tab.mode === "objects") tab.objectBrowser = { ...tab.objectBrowser, schema }; + scheduleAgentRuntimeSync(); } function updateConnection(id: string, connectionId: string, database = "") { @@ -245,6 +262,7 @@ export const useQueryStore = defineStore("query", () => { tab.resultSortedSql = undefined; clearExplain(tab); tab.tableMeta = undefined; + scheduleAgentRuntimeSync(); } function setTableMeta(id: string, meta: NonNullable) { @@ -410,6 +428,7 @@ export const useQueryStore = defineStore("query", () => { console.info("[DBX][executeTabSql:metadata:start]", { traceId, elapsed: elapsed() }); await analyzeQueryMetadata(current, current.resultBaseSql); console.info("[DBX][executeTabSql:metadata:done]", { traceId, elapsed: elapsed() }); + scheduleAgentRuntimeSync(); } else { console.warn("[DBX][executeTabSql:stale-result]", { traceId, @@ -428,6 +447,7 @@ export const useQueryStore = defineStore("query", () => { if (current.mode !== "data") current.tableMeta = undefined; current.resultBaseSql = options?.resultBaseSql ?? sql; current.resultSortedSql = options?.resultSortedSql; + scheduleAgentRuntimeSync(); } } finally { const current = tabs.value.find((t) => t.id === id); @@ -538,6 +558,7 @@ export const useQueryStore = defineStore("query", () => { if (!tab?.results || index < 0 || index >= tab.results.length) return; tab.activeResultIndex = index; tab.result = tab.results[index]; + scheduleAgentRuntimeSync(); } function trimResultCache() { diff --git a/tests/agentRuntimeSnapshot.test.ts b/tests/agentRuntimeSnapshot.test.ts new file mode 100644 index 000000000..3b02ddf16 --- /dev/null +++ b/tests/agentRuntimeSnapshot.test.ts @@ -0,0 +1,80 @@ +import { strict as assert } from "node:assert"; +import test from "node:test"; +import { buildAgentRuntimeSnapshot } from "../src/lib/agentRuntimeSnapshot.ts"; + +test("builds runtime snapshot from active tab, selection, and limited result rows", () => { + const snapshot = buildAgentRuntimeSnapshot({ + tabs: [ + { + id: "tab-1", + title: "Orders", + connectionId: "conn-1", + database: "sales", + schema: "public", + sql: "select * from orders", + result: { + columns: ["id"], + rows: [[1], [2], [3]], + affected_rows: 0, + execution_time_ms: 12, + }, + isExecuting: false, + mode: "query", + }, + ], + activeTabId: "tab-1", + getConnection: (id) => (id === "conn-1" ? { name: "Prod Sales" } : undefined), + selectedSql: "select *", + selection: { + type: "grid-cells", + data: { columns: ["id"], rows: [[1]] }, + }, + resultSampleLimit: 2, + }); + + assert.deepEqual(snapshot, { + activeConnectionId: "conn-1", + activeConnectionName: "Prod Sales", + database: "sales", + schema: "public", + activeTabId: "tab-1", + activeTabTitle: "Orders", + sql: "select * from orders", + selectedSql: "select *", + selection: { + type: "grid-cells", + data: { columns: ["id"], rows: [[1]] }, + }, + result: { + columns: ["id"], + rows: [[1], [2]], + truncated: true, + executionTimeMs: 12, + sampleLimit: 2, + }, + }); +}); + +test("omits empty selected SQL and marks empty UI selection as none", () => { + const snapshot = buildAgentRuntimeSnapshot({ + tabs: [ + { + id: "tab-2", + title: "Scratch", + connectionId: "conn-2", + database: "", + sql: "", + isExecuting: false, + mode: "query", + }, + ], + activeTabId: "tab-2", + getConnection: () => ({ name: "Local" }), + selectedSql: " ", + resultSampleLimit: 50, + }); + + assert.equal(snapshot.selectedSql, undefined); + assert.deepEqual(snapshot.selection, { type: "none" }); + assert.equal(snapshot.result, undefined); +}); From 6518061008899d1c764f10f544bd8cc7a4a826f7 Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 21:36:56 +0800 Subject: [PATCH 32/40] =?UTF-8?q?fix(agent-runtime):=20=E5=90=AF=E5=8A=A8?= =?UTF-8?q?=E6=81=A2=E5=A4=8D=E5=90=8E=E5=90=8C=E6=AD=A5=E8=BF=90=E8=A1=8C?= =?UTF-8?q?=E6=97=B6=E5=BF=AB=E7=85=A7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/App.vue | 18 +++++++++--------- src/lib/appStartup.ts | 13 +++++++++++++ tests/agentRuntimeSnapshot.test.ts | 22 ++++++++++++++++++++++ 3 files changed, 44 insertions(+), 9 deletions(-) create mode 100644 src/lib/appStartup.ts diff --git a/src/App.vue b/src/App.vue index d849be045..4db821941 100644 --- a/src/App.vue +++ b/src/App.vue @@ -38,6 +38,7 @@ import { isCloseTabShortcut, isExecuteSqlShortcut } from "@/lib/keyboardShortcut import { isPreviewTab } from "@/lib/tabPresentation"; import { SQL_FILE_UNSUPPORTED_TYPES } from "@/lib/databaseCapabilities"; import { classifyAiSqlExecution } from "@/lib/aiSqlExecutionPolicy"; +import { restoreStartupAgentRuntime } from "@/lib/appStartup"; import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; @@ -448,15 +449,14 @@ function onLoginSuccess() { } function initApp() { - savedSqlStore - .initFromStorage() - .then(() => connectionStore.initFromDisk()) - .then(() => { - reconnectRestoredTabs(); - }) - .catch((e: any) => { - toast(t("connection.loadFailed", { message: e?.message || String(e) }), 5000); - }); + restoreStartupAgentRuntime({ + initSavedSql: () => savedSqlStore.initFromStorage(), + initConnections: () => connectionStore.initFromDisk(), + reconnectRestoredTabs, + scheduleSync: () => agentRuntimeStore.scheduleSync(), + }).catch((e: any) => { + toast(t("connection.loadFailed", { message: e?.message || String(e) }), 5000); + }); settingsStore.initAiConfig(); } diff --git a/src/lib/appStartup.ts b/src/lib/appStartup.ts new file mode 100644 index 000000000..6545419cb --- /dev/null +++ b/src/lib/appStartup.ts @@ -0,0 +1,13 @@ +export interface RestoreStartupAgentRuntimeOptions { + initSavedSql: () => Promise; + initConnections: () => Promise; + reconnectRestoredTabs: () => Promise | unknown; + scheduleSync: () => void; +} + +export async function restoreStartupAgentRuntime(options: RestoreStartupAgentRuntimeOptions) { + await options.initSavedSql(); + await options.initConnections(); + await options.reconnectRestoredTabs(); + options.scheduleSync(); +} diff --git a/tests/agentRuntimeSnapshot.test.ts b/tests/agentRuntimeSnapshot.test.ts index 3b02ddf16..4a99ababa 100644 --- a/tests/agentRuntimeSnapshot.test.ts +++ b/tests/agentRuntimeSnapshot.test.ts @@ -1,6 +1,7 @@ import { strict as assert } from "node:assert"; import test from "node:test"; import { buildAgentRuntimeSnapshot } from "../src/lib/agentRuntimeSnapshot.ts"; +import { restoreStartupAgentRuntime } from "../src/lib/appStartup.ts"; test("builds runtime snapshot from active tab, selection, and limited result rows", () => { const snapshot = buildAgentRuntimeSnapshot({ @@ -78,3 +79,24 @@ test("omits empty selected SQL and marks empty UI selection as none", () => { assert.deepEqual(snapshot.selection, { type: "none" }); assert.equal(snapshot.result, undefined); }); + +test("schedules startup runtime sync only after connections and restored tabs are ready", async () => { + const calls: string[] = []; + + await restoreStartupAgentRuntime({ + initSavedSql: async () => { + calls.push("saved-sql"); + }, + initConnections: async () => { + calls.push("connections"); + }, + reconnectRestoredTabs: async () => { + calls.push("reconnect"); + }, + scheduleSync: () => { + calls.push("schedule"); + }, + }); + + assert.deepEqual(calls, ["saved-sql", "connections", "reconnect", "schedule"]); +}); From 780311a5426b1e80284140443b94b02e2abf8929 Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 21:53:11 +0800 Subject: [PATCH 33/40] =?UTF-8?q?feat(agent-runtime):=20=E6=8E=A5=E5=85=A5?= =?UTF-8?q?=20handoff=20=E5=AE=A1=E9=98=85=E7=8A=B6=E6=80=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 支持 handoff queued/shown 状态加载、mark shown 和 reject 持久化 - 为 Tauri/API/Pinia store 接入 handoff 审阅状态流转 - 新增 AgentHandoffDialog 审阅 UI，拒绝操作不执行 SQL --- crates/dbx-cli/src/commands.rs | 77 +++++++++++++++++--- crates/dbx-cli/src/main.rs | 5 +- crates/dbx-cli/src/runtime_client.rs | 34 ++------- crates/dbx-core/src/cli.rs | 8 +- crates/dbx-core/src/db/sqlite.rs | 6 +- crates/dbx-core/src/storage.rs | 74 +++++++++++++++---- src-tauri/src/commands/agent_runtime.rs | 81 +++++++++++++++++++-- src-tauri/src/commands/mod.rs | 2 +- src-tauri/src/lib.rs | 2 + src/components/agent/AgentHandoffDialog.vue | 78 ++++++++++++++++++++ src/components/layout/AppDialogs.vue | 2 + src/lib/agentHandoff.ts | 43 +++++++++++ src/lib/api.ts | 2 + src/lib/http.ts | 11 ++- src/lib/tauri.ts | 11 ++- src/stores/agentRuntimeStore.ts | 59 ++++++++++++++- tests/agentHandoff.test.ts | 57 +++++++++++++++ 17 files changed, 484 insertions(+), 68 deletions(-) create mode 100644 src/components/agent/AgentHandoffDialog.vue create mode 100644 src/lib/agentHandoff.ts create mode 100644 tests/agentHandoff.test.ts diff --git a/crates/dbx-cli/src/commands.rs b/crates/dbx-cli/src/commands.rs index 349ff0e2d..52e2324d9 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 f34308a75..9a5a19356 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 81dda04c8..effa954a0 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 534c4e1f4..6296e925a 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 a82651323..f2d668820 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 347caec01..395007291 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 8c1d093b2..d88a2ef28 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 2040e4ca0..69f61492d 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 7ce3f2cda..cca1d6a8e 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 000000000..610cd143c --- /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 0061e9d12..159fc25f0 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(