From d356135a5b898ba1f110e102ab4beb334141bf7e Mon Sep 17 00:00:00 2001 From: blankll Date: Thu, 22 Oct 2026 14:49:21 +0800 Subject: [PATCH 1/6] feat: unify AI permissions across connections and Data Studio --- DESIGN.md | 4 + docs/AI_PERMISSIONS.md | 15 + src-tauri/Cargo.lock | 3 +- src-tauri/Cargo.toml | 4 +- src-tauri/src/agent/executor.rs | 46 ++ src-tauri/src/agent_adapters.rs | 43 +- src-tauri/src/capabilities/sql_write.rs | 4 +- src-tauri/src/lib.rs | 2 + src-tauri/src/mcp_bridge.rs | 291 ++++++++-- src/__tests__/aiPermissions.test.ts | 171 ++++++ src/common/aiPermissions.ts | 14 + src/common/index.ts | 2 + .../ConnectionPermissionEditor.vue | 134 +++++ .../DefaultPermissionEditor.vue | 133 +++++ src/components/ai-permissions/index.ts | 2 + .../connections/ServerFormDialog.vue | 6 + src/composables/agentRuntime.ts | 4 +- src/composables/useDataStudioChatAgent.ts | 2 +- src/composables/useSidebarChatAgent.ts | 1 - src/datasources/agentApi.ts | 1 + src/lang/enUS.ts | 45 +- src/lang/zhCN.ts | 41 ++ src/pages/DataStudioPage.vue | 2 +- src/store/connectionStore.ts | 23 +- src/store/dataStudioStore.ts | 11 +- src/types/chat.ts | 3 +- .../components/modify-source-modal.vue | 402 +++++-------- src/views/setting/mcp-bridge.vue | 547 +++--------------- 28 files changed, 1151 insertions(+), 805 deletions(-) create mode 100644 docs/AI_PERMISSIONS.md create mode 100644 src/__tests__/aiPermissions.test.ts create mode 100644 src/common/aiPermissions.ts create mode 100644 src/components/ai-permissions/ConnectionPermissionEditor.vue create mode 100644 src/components/ai-permissions/DefaultPermissionEditor.vue create mode 100644 src/components/ai-permissions/index.ts diff --git a/DESIGN.md b/DESIGN.md index b9e2af0e..966c699d 100644 --- a/DESIGN.md +++ b/DESIGN.md @@ -328,3 +328,7 @@ Check each affected surface. Documentation-only changes need source, link, and i - Theme preview tiles scope the same light/dark tokens instead of maintaining a second miniature palette. - Automated contrast checks cover text on workspace, chrome, raised, selection, and semantic feedback surfaces, plus foregrounds on solid actions. Editor theme tests cover light/dark switching and token inheritance. - Continue to verify real content, narrow windows, keyboard focus, system appearance changes, and native window controls on the target operating systems. Theme tests do not replace native-platform interaction checks. + +### AI permission controls + +Connection dialogs expose a collapsed AI permissions section shared by Data Studio and external MCP. Default to system inheritance, show effective allowed actions, and reveal per-action Inherit / Allow / Deny under Customize. The system settings page edits action defaults, without a separate connection allowlist or MCP exposure switch. Preserve shared Select, Label, Button, neutral surfaces, localized English/Chinese copy, keyboard focus, and loading/error/retry states. See [AI permissions](docs/AI_PERMISSIONS.md) for enforcement and migration behavior. diff --git a/docs/AI_PERMISSIONS.md b/docs/AI_PERMISSIONS.md new file mode 100644 index 00000000..e12b65b5 --- /dev/null +++ b/docs/AI_PERMISSIONS.md @@ -0,0 +1,15 @@ +# AI permissions + +SQL database and document database connections share one AI permission chain across built-in Data Studio and external MCP clients. This policy does not grant database privileges or change manual database access. + +For each action (read, create data, update data, delete data, manage structure/administration), resolve an explicit connection allow/deny first, otherwise use the system default, otherwise deny. Omitted connection decisions inherit; explicit false is a denial. Presets are shortcuts for these individual decisions. + +Data Studio only operates on attached connections. Inherited attachment permissions use the live connection policy, rather than a permission snapshot. Custom conversation restrictions can narrow access, but cannot grant an operation denied by the connection chain. Ask/Auto controls confirmation and never grants CRUD permissions. Backend checks run again before execution, including retries. + +An empty set of allowed actions denies access. There is no separate MCP exposure setting or connection allowlist. Connection metadata returned to AI clients excludes inaccessible connections; Data Studio discovery is additionally limited to its attached connections. The MCP server's availability and auto-start preference remain service settings. + +Read & modify data includes data deletion but excludes structure and administration. Mixed operations require every affected permission: SQL writes inspect statement tokens, bulk requests inspect every action, and aggregation pipelines that write require write and management access. Risk controls confirmation independently from action permissions. + +Legacy allowlists migrate to explicit connection decisions preserving exclusions. When an allowlist existed, new connections default to no access. Existing per-connection restrictions are preserved. Corrupt or unreadable policy files deny execution rather than falling back to permissive settings. Policy writes are atomic and take effect for subsequent calls without restarting the bridge. + +Both applications pin the shared permission implementation to an exact Git revision in their Cargo manifests and lockfiles so builds use the same policy behavior. diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index 265c8880..434301e3 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -1554,7 +1554,7 @@ checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" [[package]] name = "data-studio-agent" version = "0.1.9" -source = "git+https://github.com/geek-fun/data-studio-agent.git?tag=v0.1.9#7e6610bd421d5acbadf64558aa954bc27b175e38" +source = "git+https://github.com/geek-fun/data-studio-agent.git?rev=c36fb2b16a4031053c3990dcd4d4b8e84d4d597f#c36fb2b16a4031053c3990dcd4d4b8e84d4d597f" dependencies = [ "async-openai", "async-trait", @@ -1564,6 +1564,7 @@ dependencies = [ "rusqlite", "serde", "serde_json", + "sqlparser", "tiktoken-rs", "tokio", "uuid", diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index ca52287a..be3967d2 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -87,8 +87,8 @@ http = "1" log = "0.4" futures = "0.3" rand = "0.8" -# data-studio-agent: pinned to latest release tag. -data-studio-agent = { git = "https://github.com/geek-fun/data-studio-agent.git", tag = "v0.1.9" } +# Pin the shared AI permission chain to a reproducible revision. +data-studio-agent = { git = "https://github.com/geek-fun/data-studio-agent.git", rev = "c36fb2b16a4031053c3990dcd4d4b8e84d4d597f" } # Archive extraction (JRE downloads) flate2 = "1.0" diff --git a/src-tauri/src/agent/executor.rs b/src-tauri/src/agent/executor.rs index 68512ed5..73d2afb4 100644 --- a/src-tauri/src/agent/executor.rs +++ b/src-tauri/src/agent/executor.rs @@ -24,6 +24,25 @@ impl ToolExecutor for SqlKitToolExecutor { connection_config: &Value, ) -> Result { let start = std::time::Instant::now(); + let app = crate::APP_HANDLE.get().ok_or("App not initialized")?; + let policy = crate::mcp_bridge::load_permission_policy(app)?; + let cap = data_studio_agent::capabilities::registry::registry() + .get(tool_name) + .ok_or("Unknown capability")?; + let id = connection_config + .get("connectionId") + .and_then(Value::as_str); + let restriction = connection_config + .get("permissions") + .filter(|v| !v.is_null()) + .map(|v| { + serde_json::from_value::< + data_studio_agent::capabilities::permissions::PermissionRules, + >(v.clone()) + }) + .transpose() + .map_err(|e| e.to_string())?; + policy.authorize(cap, arguments, id, restriction.as_ref())?; let conn_opt = if connection_config.is_null() { None @@ -38,6 +57,33 @@ impl ToolExecutor for SqlKitToolExecutor { ) .await?; + let raw = if tool_name.ends_with("__list_connections") { + let mut result: Value = serde_json::from_str(&raw).map_err(|e| e.to_string())?; + if let Some(connections) = result.get("connections").and_then(Value::as_array) { + let filtered = connections + .iter() + .filter(|connection| { + let id = match &connection["id"] { + Value::String(id) => id.clone(), + Value::Number(id) => id.to_string(), + _ => return false, + }; + let in_scope = connection_config + .get("__allowed_connection_ids") + .and_then(Value::as_array) + .is_some_and(|ids| { + ids.iter().any(|allowed| allowed.as_str() == Some(&id)) + }); + in_scope && policy.effective_permissions(&id).any_allowed() + }) + .cloned() + .collect::>(); + result["connections"] = serde_json::json!(filtered); + } + result.to_string() + } else { + raw + }; let duration_ms = start.elapsed().as_millis() as u64; let (full_result, truncated) = char_truncate(&raw, TOOL_ENVELOPE_MAX_CHARS); let (summary, _) = char_truncate(&full_result, TOOL_ENVELOPE_SUMMARY_CHARS); diff --git a/src-tauri/src/agent_adapters.rs b/src-tauri/src/agent_adapters.rs index 739e483f..58192952 100644 --- a/src-tauri/src/agent_adapters.rs +++ b/src-tauri/src/agent_adapters.rs @@ -7,7 +7,6 @@ use std::collections::HashMap; use std::sync::Arc; -use crate::mcp_bridge::McpConfig; use data_studio_agent as lib; use data_studio_agent::storage; use data_studio_agent::traits::{CancelMap, ConfirmMap, EventEmitter}; @@ -69,29 +68,29 @@ pub async fn run_agent_loop( .unwrap_or(false) }; - let app_data_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + let policy_app = app.clone(); + let permission_connections = connections.clone(); let should_deny = move |tool_name: &str, conn_id: Option<&str>| -> bool { - // Reload per call so permission changes made in Settings apply to the - // next tool call instead of being snapshotted for the whole run. - let policy = McpConfig::load(&app_data_dir).policy; - match data_studio_agent::capabilities::registry::registry().get(tool_name) { - Some(cap) => { - // Allowlist set but no connection specified: a DB capability would - // fall back to the default connection, bypassing the allowlist. - // AppLocal tools (list_connections etc.) need no connection. - if !policy.allowed_connection_ids.is_empty() - && conn_id.is_none() - && !matches!( - cap.source_kind, - data_studio_agent::capabilities::types::SourceKind::AppLocal - ) - { - return true; - } - !policy.allows(cap.risk_level, conn_id) - } - None => false, + let Some(cap) = data_studio_agent::capabilities::registry::registry().get(tool_name) else { + return true; + }; + if matches!( + cap.source_kind, + data_studio_agent::capabilities::types::SourceKind::AppLocal + | data_studio_agent::capabilities::types::SourceKind::File + ) { + return false; } + let Some(id) = conn_id + .and_then(|id| permission_connections.get(id)) + .and_then(|cfg| cfg.get("connectionId")) + .and_then(Value::as_str) + else { + return true; + }; + crate::mcp_bridge::load_permission_policy(&policy_app) + .map(|policy| !policy.effective_permissions(id).any_allowed()) + .unwrap_or(true) }; lib::loop_runner::run_agent_loop( diff --git a/src-tauri/src/capabilities/sql_write.rs b/src-tauri/src/capabilities/sql_write.rs index 3c13f37e..0cec1d02 100644 --- a/src-tauri/src/capabilities/sql_write.rs +++ b/src-tauri/src/capabilities/sql_write.rs @@ -366,7 +366,7 @@ pub(crate) fn register_write_tools(reg: &mut CapabilityRegistry) { reg.register(Capability { name: "sqlkit__execute_delete", - description: "Execute a destructive DELETE or TRUNCATE statement. DESTRUCTIVE: permanently removes data. Requires Full Access with Confirm Destructive in Settings → MCP Bridge.", + description: "Execute a destructive DELETE or TRUNCATE statement. DESTRUCTIVE: permanently removes data. Requires Delete data permission for this connection.", handler: Arc::new(ExecuteDeleteHandler), input_schema: json!({"type": "object", "properties": { "connection_id": connection_id_schema(), @@ -381,7 +381,7 @@ pub(crate) fn register_write_tools(reg: &mut CapabilityRegistry) { reg.register(Capability { name: "sqlkit__execute_ddl", - description: "Execute a DDL statement (CREATE, ALTER, DROP, and other schema changes). DESTRUCTIVE for DROP/ALTER: permanently changes schema. Requires Full Access with Confirm Destructive in Settings → MCP Bridge.", + description: "Execute a DDL statement (CREATE, ALTER, DROP, and other schema changes). Changes database structure. Requires Manage structure & administration permission for this connection.", handler: Arc::new(ExecuteDdlHandler), input_schema: json!({"type": "object", "properties": { "connection_id": connection_id_schema(), diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index a1298c46..1228e729 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -268,6 +268,8 @@ pub fn run() { Ok(()) }) .invoke_handler(tauri::generate_handler![ + crate::mcp_bridge::get_ai_permissions, + crate::mcp_bridge::save_ai_permissions, // SSH profile management + host-key store crate::ssh::commands::list_ssh_profiles, crate::ssh::commands::save_ssh_profile, diff --git a/src-tauri/src/mcp_bridge.rs b/src-tauri/src/mcp_bridge.rs index 6caef5b7..3b566fed 100644 --- a/src-tauri/src/mcp_bridge.rs +++ b/src-tauri/src/mcp_bridge.rs @@ -24,9 +24,6 @@ use tokio::sync::oneshot; pub struct McpServerHandle { pub shutdown_tx: Mutex>>, pub server_task: Mutex>>, - /// Shared policy published into the running bridge. `save_mcp_policy` - /// hot-swaps through this without restarting the server. - pub policy: Arc>, } impl McpServerHandle { @@ -34,7 +31,6 @@ impl McpServerHandle { Self { shutdown_tx: Mutex::new(None), server_task: Mutex::new(None), - policy: Arc::new(tokio::sync::RwLock::new(McpPolicy::default())), } } } @@ -78,7 +74,13 @@ impl McpConfig { pub fn save(&self, app_data_dir: &Path) -> Result<(), String> { let path = app_data_dir.join("mcp-config.json"); let json = serde_json::to_string_pretty(self).map_err(|e| e.to_string())?; - std::fs::write(&path, &json).map_err(|e| format!("Failed to write mcp-config.json: {}", e)) + let temporary = app_data_dir.join(format!("mcp-config-{}.tmp", uuid::Uuid::new_v4())); + std::fs::write(&temporary, &json) + .map_err(|e| format!("Failed to write MCP config: {e}"))?; + std::fs::rename(&temporary, &path).map_err(|e| { + let _ = std::fs::remove_file(&temporary); + format!("Failed to save MCP config: {e}") + }) } } @@ -92,6 +94,84 @@ impl Default for McpConfig { } } +static AI_POLICY_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + +pub(crate) fn load_permission_policy( + app: &tauri::AppHandle, +) -> Result { + let _guard = AI_POLICY_LOCK.lock().map_err(|e| e.to_string())?; + load_permission_policy_inner(app) +} + +fn read_permission_config(dir: &Path) -> Result { + Ok(match std::fs::read_to_string(dir.join("mcp-config.json")) { + Ok(raw) => { + serde_json::from_str(&raw).map_err(|e| format!("Cannot load AI permissions: {e}"))? + } + Err(error) if error.kind() == std::io::ErrorKind::NotFound => McpConfig::default(), + Err(error) => return Err(format!("Cannot load AI permissions: {error}")), + }) +} + +fn load_permission_policy_inner( + app: &tauri::AppHandle, +) -> Result { + let dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + let mut config = read_permission_config(&dir)?; + if config.policy.permission_defaults.is_none() { + let store = app.store(".store.dat").map_err(|e| e.to_string())?; + let ids = store + .get("connections") + .and_then(|v| v.as_array().cloned()) + .unwrap_or_default() + .iter() + .filter_map(|c| { + c.get("id").and_then(|id| match id { + Value::String(id) => Some(id.clone()), + Value::Number(id) => Some(id.to_string()), + _ => None, + }) + }) + .collect::>(); + config.policy.migrate_permissions(&ids); + std::fs::create_dir_all(&dir).map_err(|e| e.to_string())?; + config.save(&dir)?; + } + Ok(config.policy) +} + +#[tauri::command] +pub fn get_ai_permissions(app: AppHandle) -> Result { + load_permission_policy(&app) +} + +#[tauri::command] +pub fn save_ai_permissions( + defaults: Option, + connection_id: Option, + permissions: Option, + app: AppHandle, +) -> Result { + let _guard = AI_POLICY_LOCK.lock().map_err(|e| e.to_string())?; + let mut policy = load_permission_policy_inner(&app)?; + if let Some(defaults) = defaults { + policy.permission_defaults = Some(defaults.resolve(&Default::default())); + } + if let Some(id) = connection_id { + let rules = permissions.unwrap_or_default(); + if rules == Default::default() { + policy.connection_permissions.remove(&id); + } else { + policy.connection_permissions.insert(id, rules); + } + } + let dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + let mut config = read_permission_config(&dir)?; + config.policy = policy.clone(); + config.save(&dir)?; + Ok(policy) +} + // --------------------------------------------------------------------------- // Request / response types // --------------------------------------------------------------------------- @@ -138,7 +218,6 @@ struct BridgeState { handle: AppHandle, app_name: &'static str, app_data_dir: PathBuf, - policy: Arc>, } // --------------------------------------------------------------------------- @@ -152,7 +231,7 @@ fn tools_payload(policy: &McpPolicy) -> Value { let tools: Vec = caps .iter() - .filter(|cap| check_policy(cap, policy, None).is_ok()) + .filter(|cap| policy.advertises(cap)) .map(|cap| { // Neon-style policy notice: tells the client (LLM) which // capability classes are gated and how to lift the gate, @@ -201,7 +280,16 @@ async fn handle_tools(State(state): State>) -> Response { ) .into_response(); } - let policy = state.policy.read().await; + let policy = match load_permission_policy(&state.handle) { + Ok(policy) => policy, + Err(e) => { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({"message": e})), + ) + .into_response() + } + }; let mut result = tools_payload(&policy); result["connections"] = list_allowed_connections(&policy); Json(result).into_response() @@ -256,6 +344,9 @@ fn db_type_str(c: &serde_json::Value) -> String { /// Saved connections restricted to the policy allowlist, so MCP clients only /// learn about connections they are permitted to reach. fn list_allowed_connections(policy: &McpPolicy) -> Value { + if policy.permission_defaults.is_some() { + return filter_ai_connections(list_connections(), policy); + } if policy.allowed_connection_ids.is_empty() { return list_connections(); } @@ -303,8 +394,15 @@ fn check_policy( /// Deny reason for a connection-less invocation when an allowlist is set. /// AppLocal capabilities need no connection, so they stay exempt. -fn allowlist_missing_connection_reason(cap: &Capability, policy: &McpPolicy) -> Option { - if !policy.allowed_connection_ids.is_empty() && !matches!(cap.source_kind, SourceKind::AppLocal) +fn allowlist_missing_connection_reason( + cap: &Capability, + policy: &McpPolicy, + connection_id: Option<&str>, +) -> Option { + if connection_id.is_none() + && policy.permission_defaults.is_none() + && !policy.allowed_connection_ids.is_empty() + && !matches!(cap.source_kind, SourceKind::AppLocal) { return Some( "connection allowlist is set — supply a connection_id from the allowlist in Settings → MCP Bridge".into(), @@ -322,13 +420,35 @@ async fn handle_invoke( } // Snapshot so an in-flight request keeps its policy if the user hot-swaps // permissions mid-request. - let policy = state.policy.read().await.clone(); + let policy = match load_permission_policy(&state.handle) { + Ok(policy) => policy, + Err(e) => return Json(InvokeResponse::error(500, e)), + }; Json(invoke_with_policy(&policy, payload).await) } // Split from handle_invoke so tests can run it without a Tauri runtime — // BridgeState holds a Wry AppHandle, which tauri::test mocks cannot provide. async fn invoke_with_policy(policy: &McpPolicy, payload: InvokeRequest) -> InvokeResponse { + if let (Some(envelope), Some(argument)) = ( + payload.connection_id.as_deref(), + payload.args.get("connection_id").and_then(Value::as_str), + ) { + if envelope != argument { + return InvokeResponse::error(400, "Conflicting connection identifiers".into()); + } + } + let payload = InvokeRequest { + connection_id: payload.connection_id.or_else(|| { + payload + .args + .get("connection_id") + .and_then(Value::as_str) + .map(str::to_owned) + }), + ..payload + }; + let cap = match registry::registry().get(&payload.name) { Some(c) => c, None => return InvokeResponse::error(404, format!("Unknown capability: {}", payload.name)), @@ -337,7 +457,9 @@ async fn invoke_with_policy(policy: &McpPolicy, payload: InvokeRequest) -> Invok // Allowlist set but no connection specified: a DB capability would fall // back to the default connection, bypassing the allowlist. Mirrors the // guard in the built-in agent loop (agent_adapters.rs). - if let Some(reason) = allowlist_missing_connection_reason(cap, policy) { + if let Some(reason) = + allowlist_missing_connection_reason(cap, policy, payload.connection_id.as_deref()) + { return InvokeResponse::error( 403, format!( @@ -347,7 +469,13 @@ async fn invoke_with_policy(policy: &McpPolicy, payload: InvokeRequest) -> Invok ); } - if let Err(msg) = check_policy(cap, policy, payload.connection_id.as_deref()) { + if let Err(msg) = if policy.permission_defaults.is_some() { + policy + .authorize(cap, &payload.args, payload.connection_id.as_deref(), None) + .map(|_| ()) + } else { + check_policy(cap, policy, payload.connection_id.as_deref()) + } { return InvokeResponse::error(403, msg); } @@ -361,7 +489,14 @@ async fn invoke_with_policy(policy: &McpPolicy, payload: InvokeRequest) -> Invok match registry::invoke_capability_inner(&payload.name, payload.args, config).await { Ok(data) => match serde_json::from_str::(&data) { - Ok(parsed) => InvokeResponse::ok(parsed), + Ok(mut parsed) => { + if payload.name.ends_with("__list_connections") { + if let Some(connections) = parsed.get("connections") { + parsed["connections"] = filter_ai_connections(connections.clone(), policy); + } + } + InvokeResponse::ok(parsed) + } Err(_) => InvokeResponse::ok(json!({"result": data})), }, Err(msg) => InvokeResponse::error(400, msg), @@ -507,14 +642,6 @@ pub async fn start( preferred_port: u16, shutdown_rx: oneshot::Receiver<()>, ) -> Result { - // Publish the persisted policy into the shared Arc so the running bridge - // and save_mcp_policy hot-swaps stay in sync. - let policy_arc = { - let mcp_handle: tauri::State<'_, McpServerHandle> = handle.state(); - *mcp_handle.policy.write().await = McpConfig::load(&app_data_dir).policy; - mcp_handle.policy.clone() - }; - // Try preferred port first, fall back to random let port = match TcpListener::bind(format!("127.0.0.1:{}", preferred_port)).await { Ok(listener) => { @@ -522,7 +649,6 @@ pub async fn start( handle: handle.clone(), app_name: "sqlkit", app_data_dir: app_data_dir.clone(), - policy: policy_arc.clone(), }); let app = axum::Router::new() @@ -568,7 +694,6 @@ pub async fn start( handle: handle.clone(), app_name: "sqlkit", app_data_dir: app_data_dir.clone(), - policy: policy_arc.clone(), }); let app = axum::Router::new() @@ -633,7 +758,7 @@ pub async fn get_mcp_status(app: AppHandle) -> Result { .app_data_dir() .map_err(|e| format!("Failed to resolve app data dir: {}", e))?; - let config = McpConfig::load(&app_data_dir); + let config = read_permission_config(&app_data_dir)?; let running_port = get_actual_port(&app_data_dir); let status = json!({ @@ -664,14 +789,17 @@ pub async fn save_mcp_config( .map_err(|e| format!("Failed to resolve app data dir: {}", e))? .to_path_buf(); - // Load first so a None policy from older clients keeps the stored policy. - let mut config = McpConfig::load(&app_data_dir); - config.port = port; - config.auto_start = auto_start; - if let Some(p) = policy { - config.policy = p; + { + let _guard = AI_POLICY_LOCK.lock().map_err(|e| e.to_string())?; + // Load first so a None policy from older clients keeps the stored policy. + let mut config = read_permission_config(&app_data_dir)?; + config.port = port; + config.auto_start = auto_start; + if let Some(p) = policy { + config.policy = p; + } + config.save(&app_data_dir)?; } - config.save(&app_data_dir)?; // Explicit restart starts now without persisting a change to auto_start. let should_start = auto_start || restart.unwrap_or(false); @@ -716,12 +844,12 @@ pub async fn save_mcp_policy(policy: McpPolicy, app: AppHandle) -> Result = app.state(); - *server_handle.policy.write().await = policy; + { + let _guard = AI_POLICY_LOCK.lock().map_err(|e| e.to_string())?; + let mut config = read_permission_config(&app_data_dir)?; + config.policy = policy.clone(); + config.save(&app_data_dir)?; + } serde_json::to_string(&json!({"status": "ok"})).map_err(|e| e.to_string()) } @@ -1023,11 +1151,12 @@ mod tests { ..McpPolicy::default() }; // DB capability without a connection_id → denied by allowlist - assert!(allowlist_missing_connection_reason(db_cap, &allowlist).is_some()); + assert!(allowlist_missing_connection_reason(db_cap, &allowlist, None).is_some()); + assert!(allowlist_missing_connection_reason(db_cap, &allowlist, Some("conn-1")).is_none()); // Empty allowlist → no restriction - assert!(allowlist_missing_connection_reason(db_cap, &McpPolicy::default()).is_none()); + assert!(allowlist_missing_connection_reason(db_cap, &McpPolicy::default(), None).is_none()); // AppLocal capabilities never need a connection - assert!(allowlist_missing_connection_reason(app_local_cap, &allowlist).is_none()); + assert!(allowlist_missing_connection_reason(app_local_cap, &allowlist, None).is_none()); } #[test] @@ -1165,4 +1294,84 @@ mod tests { ); } } + #[test] + fn ai_permission_roundtrip_preserves_partial_decisions_and_service_settings() { + use data_studio_agent::capabilities::permissions::PermissionRules; + let dir = std::env::temp_dir().join(format!("sqlkit-ai-chain-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let mut policy = McpPolicy { + permission_defaults: Some(PermissionRules::explicit(true, false, false, false, false)), + ..Default::default() + }; + policy.connection_permissions.insert( + "42".into(), + PermissionRules { + update: Some(true), + delete: Some(false), + ..Default::default() + }, + ); + let config = McpConfig { + port: Some(9432), + auto_start: false, + policy: policy.clone(), + }; + config.save(&dir).unwrap(); + let loaded = read_permission_config(&dir).unwrap(); + assert_eq!(loaded.policy, policy); + assert_eq!(loaded.port, Some(9432)); + assert!(!loaded.auto_start); + assert_eq!(loaded.policy.effective_permissions("42").update, Some(true)); + assert_eq!( + loaded.policy.effective_permissions("42").delete, + Some(false) + ); + std::fs::write(dir.join("mcp-config.json"), "invalid").unwrap(); + assert!(read_permission_config(&dir).is_err()); + std::fs::remove_dir_all(dir).unwrap(); + } + + #[test] + fn ai_connection_discovery_filters_numeric_and_string_ids_by_effective_permissions() { + use data_studio_agent::capabilities::permissions::PermissionRules; + let mut policy = McpPolicy { + permission_defaults: Some(PermissionRules::explicit(false, false, false, false, false)), + ..Default::default() + }; + policy.connection_permissions.insert( + "42".into(), + PermissionRules { + read: Some(true), + ..Default::default() + }, + ); + let filtered = + filter_ai_connections(json!([{"id":42},{"id":"other"},{"id":null}]), &policy); + assert_eq!(filtered, json!([{"id":42}])); + policy.connection_permissions.insert( + "42".into(), + PermissionRules { + read: Some(false), + ..Default::default() + }, + ); + assert_eq!(filter_ai_connections(filtered, &policy), json!([])); + } +} + +fn filter_ai_connections(connections: Value, policy: &McpPolicy) -> Value { + json!(connections + .as_array() + .into_iter() + .flatten() + .filter(|connection| { + let id = match &connection["id"] { + Value::String(id) => id.clone(), + Value::Number(id) => id.to_string(), + _ => return false, + }; + policy.effective_permissions(&id).any_allowed() + }) + .cloned() + .collect::>()) } diff --git a/src/__tests__/aiPermissions.test.ts b/src/__tests__/aiPermissions.test.ts new file mode 100644 index 00000000..af4f6e6a --- /dev/null +++ b/src/__tests__/aiPermissions.test.ts @@ -0,0 +1,171 @@ +/** @jest-environment jsdom */ +import { invoke } from '@tauri-apps/api/core' +import { flushPromises, mount } from '@vue/test-utils' +import { defineComponent } from 'vue' +import { AI_ACTIONS, aiPreset, resolveAiPermissions } from '@/common/aiPermissions' +import ConnectionPermissionEditor from '@/components/ai-permissions/ConnectionPermissionEditor.vue' +import DefaultPermissionEditor from '@/components/ai-permissions/DefaultPermissionEditor.vue' +import SourcePermissionModal from '@/views/data-studio/components/modify-source-modal.vue' + +jest.mock('@tauri-apps/api/core', () => ({ invoke: jest.fn() })) +jest.mock('@/common', () => jest.requireActual('@/common/aiPermissions')) +jest.mock('vue-i18n', () => ({ useI18n: () => ({ t: (key: string) => key }) })) +jest.mock('@/components/ui/button', () => ({ Button: { template: '' } })) +jest.mock('@/components/ui/label', () => ({ Label: { template: '' } })) +jest.mock('@/components/ui/select', () => ({ + Select: { props: ['modelValue'], emits: ['update:modelValue'], template: '' }, + SelectTrigger: { template: '' }, + SelectValue: { template: '' }, + SelectContent: { template: '' }, + SelectItem: { props: ['value', 'disabled'], template: '' }, +})) + +jest.mock('@/components/ui/dialog', () => Object.fromEntries( + ['Dialog', 'DialogContent', 'DialogDescription', 'DialogFooter', 'DialogHeader', 'DialogTitle'] + .map(name => [name, { template: '
' }]), +)) +jest.mock('@/components/ui/checkbox', () => ({ + Checkbox: { + props: ['checked', 'disabled'], + emits: ['update:checked'], + template: '', + }, +})) + +const mockStore = { + activeSession: { id: 'session', sources: [{ sourceId: 'source', alias: 'Example', permissionsMode: 'custom', permissions: { read: true, update: true } }] }, + attachedSources: [{ sourceId: 'source', kind: 'database', connectionId: 'dev' }], + updateSessionSourcePermissions: jest.fn(), + updateSessionSourceMode: jest.fn(), + detachSourceFromSession: jest.fn(), +} + +jest.mock('@/store/dataStudioStore', () => ({ useDataStudioStore: () => mockStore })) + +const mockInvoke = jest.mocked(invoke) +const policy = { permission_defaults: aiPreset('read'), connection_permissions: { dev: { update: true, delete: false } } } +const openEditor = async (id?: string) => { + const wrapper = mount(defineComponent({ + components: { ConnectionPermissionEditor }, + data: () => ({ rules: {}, id }), + template: '', + })) + await flushPromises() + await wrapper.get('summary').trigger('click') + return wrapper +} + +describe('aI permission chain', () => { + beforeEach(() => { + mockStore.updateSessionSourcePermissions.mockReset() + mockStore.detachSourceFromSession.mockReset() + mockInvoke.mockReset() + mockInvoke.mockResolvedValue(policy) + }) + + it('resolves each action with explicit false distinct from inheritance', () => { + expect(resolveAiPermissions(aiPreset('full'), { update: false })).toEqual({ read: true, create: true, update: false, delete: true, manage: true }) + expect(resolveAiPermissions({}, { read: true })).toEqual(aiPreset('read')) + expect(aiPreset('inherit')).toEqual({}) + expect(AI_ACTIONS.every(action => aiPreset('none')[action] === false)).toBe(true) + }) + + it('loads saved decisions and displays effective access from the whole chain', async () => { + const wrapper = await openEditor('dev') + expect(wrapper.vm.rules).toEqual({ update: true, delete: false }) + expect(wrapper.text()).toContain('aiPermissions.effective: aiPermissions.read, aiPermissions.update') + await wrapper.get('button').trigger('click') + expect(wrapper.text()).toContain('aiPermissions.fromConnection') + expect(wrapper.text()).toContain('aiPermissions.fromSystem') + }) + + it('supports inherit, explicit allow and explicit deny independently', async () => { + const wrapper = await openEditor('dev') + await wrapper.get('button').trigger('click') + const actionSelect = () => wrapper.findAll('select')[3] + await actionSelect().setValue('deny') + expect(wrapper.vm.rules).toEqual({ update: false, delete: false }) + await actionSelect().setValue('inherit') + expect(wrapper.vm.rules).toEqual({ delete: false }) + await actionSelect().setValue('allow') + expect(wrapper.vm.rules).toEqual({ update: true, delete: false }) + }) + + it('keeps custom decisions when opening the custom preset', async () => { + const wrapper = await openEditor('dev') + await wrapper.get('select').setValue('custom') + expect(wrapper.vm.rules).toEqual({ update: true, delete: false }) + expect(wrapper.findAll('select')).toHaveLength(6) + }) + + it('new connections inherit and presets preserve explicit decisions', async () => { + const wrapper = await openEditor() + expect(wrapper.vm.rules).toEqual({}) + await wrapper.get('select').setValue('none') + expect(wrapper.vm.rules).toEqual(aiPreset('none')) + await wrapper.get('select').setValue('inherit') + expect(wrapper.vm.rules).toEqual({}) + }) + + it('blocks readiness after load failure and recovers through Retry', async () => { + mockInvoke.mockRejectedValueOnce(new Error('Unavailable')) + const wrapper = await openEditor('dev') + const editor = wrapper.findComponent(ConnectionPermissionEditor) + expect((editor.vm as unknown as { ready: boolean }).ready).toBe(false) + expect(wrapper.text()).toContain('aiPermissions.loadFailed') + await wrapper.get('button').trigger('click') + await flushPromises() + expect((editor.vm as unknown as { ready: boolean }).ready).toBe(true) + }) + + it('retries failed default saves without discarding the edited permissions', async () => { + const wrapper = mount(DefaultPermissionEditor) + await flushPromises() + await wrapper.get('select').setValue('full') + mockInvoke.mockRejectedValueOnce(new Error('Temporarily unavailable')) + await wrapper.findAll('button').find(button => button.text() === 'aiPermissions.save')!.trigger('click') + await flushPromises() + expect(wrapper.get('select').element.value).toBe('full') + await wrapper.findAll('button').find(button => button.text() === 'aiPermissions.retry')!.trigger('click') + await flushPromises() + expect(mockInvoke).toHaveBeenLastCalledWith('save_ai_permissions', { defaults: aiPreset('full') }) + expect(wrapper.text()).toContain('aiPermissions.saved') + }) + + it('keeps conversation edits after a failed save and retries the same draft', async () => { + const wrapper = mount(SourcePermissionModal, { props: { open: true, sourceIdx: 0 } }) + await flushPromises() + expect(wrapper.get('#conversation-create').attributes('disabled')).toBeDefined() + await wrapper.get('#conversation-update').setValue(false) + mockStore.updateSessionSourcePermissions.mockRejectedValueOnce(new Error('Unavailable')) + await wrapper.findAll('button').find(button => button.text() === 'dataStudio.modifySource.saveChanges')!.trigger('click') + await flushPromises() + expect(wrapper.emitted('update:open')).toBeUndefined() + expect((wrapper.get('#conversation-update').element as HTMLInputElement).checked).toBe(false) + await wrapper.findAll('button').find(button => button.text() === 'aiPermissions.retry')!.trigger('click') + await flushPromises() + expect(mockStore.updateSessionSourcePermissions).toHaveBeenLastCalledWith('source', { read: true, create: false, update: false, delete: false, manage: false }) + expect(wrapper.emitted('update:open')).toEqual([[false]]) + }) + + it('requires a separate detach confirmation and lets the user cancel it', async () => { + const wrapper = mount(SourcePermissionModal, { props: { open: true, sourceIdx: 0 } }) + await flushPromises() + await wrapper.findAll('button').find(button => button.text() === 'dataStudio.detachSource.title')!.trigger('click') + expect(wrapper.text()).toContain('dataStudio.detachSource.message') + expect(mockStore.detachSourceFromSession).not.toHaveBeenCalled() + await wrapper.findAll('button').find(button => button.text() === 'aiPermissions.cancel')!.trigger('click') + expect(wrapper.text()).not.toContain('dataStudio.detachSource.message') + expect(wrapper.emitted('update:open')).toBeUndefined() + }) + + it('saves system defaults without overwriting connection decisions', async () => { + const wrapper = mount(DefaultPermissionEditor) + await flushPromises() + await wrapper.get('select').setValue('none') + await wrapper.findAll('button').find(button => button.text() === 'aiPermissions.save')!.trigger('click') + await flushPromises() + expect(mockInvoke).toHaveBeenLastCalledWith('save_ai_permissions', { defaults: aiPreset('none') }) + expect(wrapper.text()).toContain('aiPermissions.saved') + }) +}) diff --git a/src/common/aiPermissions.ts b/src/common/aiPermissions.ts new file mode 100644 index 00000000..98cdd95e --- /dev/null +++ b/src/common/aiPermissions.ts @@ -0,0 +1,14 @@ +import { invoke } from '@tauri-apps/api/core' + +export const AI_ACTIONS = ['read', 'create', 'update', 'delete', 'manage'] as const +export type AiAction = typeof AI_ACTIONS[number] +export type AiPermissions = Partial> +export type AiPolicy = { permission_defaults: AiPermissions, connection_permissions: Record } +export const getAiPolicy = () => invoke('get_ai_permissions') +export const saveAiPermissions = (connectionId: string, permissions: AiPermissions) => + invoke('save_ai_permissions', { connectionId, permissions }) +export const resolveAiPermissions = (defaults: AiPermissions, connection: AiPermissions = {}): AiPermissions => + Object.fromEntries(AI_ACTIONS.map(action => [action, connection[action] ?? defaults[action] ?? false])) +export const aiPreset = (value: string): AiPermissions => value === 'inherit' + ? {} + : Object.fromEntries(AI_ACTIONS.map(action => [action, value === 'full' || (value === 'read' && action === 'read') || (value === 'data' && action !== 'manage')])) diff --git a/src/common/index.ts b/src/common/index.ts index 33125bdf..01752896 100644 --- a/src/common/index.ts +++ b/src/common/index.ts @@ -4,7 +4,9 @@ export const pureObject = (obj: T): T => JSON.parse(JSON.stringify(obj)) as T export * from './agentPrompt' +export * from './aiPermissions' export * from './entitlement' export * from './sessionExpiry' export * from './sessionPhase' + export * from './sqlParser' diff --git a/src/components/ai-permissions/ConnectionPermissionEditor.vue b/src/components/ai-permissions/ConnectionPermissionEditor.vue new file mode 100644 index 00000000..e3800590 --- /dev/null +++ b/src/components/ai-permissions/ConnectionPermissionEditor.vue @@ -0,0 +1,134 @@ + + + diff --git a/src/components/ai-permissions/DefaultPermissionEditor.vue b/src/components/ai-permissions/DefaultPermissionEditor.vue new file mode 100644 index 00000000..d969c625 --- /dev/null +++ b/src/components/ai-permissions/DefaultPermissionEditor.vue @@ -0,0 +1,133 @@ + + + diff --git a/src/components/ai-permissions/index.ts b/src/components/ai-permissions/index.ts new file mode 100644 index 00000000..5f47a002 --- /dev/null +++ b/src/components/ai-permissions/index.ts @@ -0,0 +1,2 @@ +export { default as ConnectionPermissionEditor } from './ConnectionPermissionEditor.vue' +export { default as DefaultPermissionEditor } from './DefaultPermissionEditor.vue' diff --git a/src/components/connections/ServerFormDialog.vue b/src/components/connections/ServerFormDialog.vue index bb1a56c1..6fc1d4d1 100644 --- a/src/components/connections/ServerFormDialog.vue +++ b/src/components/connections/ServerFormDialog.vue @@ -7,6 +7,7 @@ import { openUrl } from '@tauri-apps/plugin-opener' import { computed, ref, watch } from 'vue' import { useI18n } from 'vue-i18n' import { shouldStripSshTunnel, UPGRADE_URL } from '@/common' +import { ConnectionPermissionEditor } from '@/components/ai-permissions' import { Badge } from '@/components/ui/badge' import { Button } from '@/components/ui/button' import { SearchableSelect } from '@/components/ui/combobox' @@ -61,6 +62,7 @@ const isOpen = computed({ }) const isEditing = computed(() => !!props.connection?.id) +const aiPermissionEditor = ref | null>(null) const formSectionTab = ref<'basic' | 'advanced'>('basic') const defaultPorts: Record = { @@ -896,6 +898,8 @@ function mapDatabaseTypeToBackend(type: DatabaseType): string { } function handleSave() { + if (aiPermissionEditor.value && !aiPermissionEditor.value.ready) + return if (!validateForm()) { return } @@ -1487,6 +1491,8 @@ function handleSave() { :errors="formErrors" /> + + { ) unlisteners.push( - await agentApi.onAgentLoopToolCall(({ session_id, tool_call_id, tool_name, arguments: args }) => { + await agentApi.onAgentLoopToolCall(({ session_id, tool_call_id, tool_name, arguments: args, risk_level }) => { const store = useDataStudioStore() store.removePreparingPlaceholder() store.setSessionActiveTool(session_id, tool_name) @@ -157,7 +157,7 @@ async function initAgentRuntime(): Promise { return const runtime = getSessionRuntime(session_id) - const riskLevel = (runtime.metadata?.[tool_name]?.riskLevel as AgentToolCall['riskLevel']) ?? 'elevated' + const riskLevel = risk_level ?? (runtime.metadata?.[tool_name]?.riskLevel as AgentToolCall['riskLevel']) ?? 'elevated' const needsConfirmation = shouldRequireConfirmation(session, tool_name, riskLevel) const denied = isDeniedByRule(session_id, tool_name) diff --git a/src/composables/useDataStudioChatAgent.ts b/src/composables/useDataStudioChatAgent.ts index 7514d997..71d5dace 100644 --- a/src/composables/useDataStudioChatAgent.ts +++ b/src/composables/useDataStudioChatAgent.ts @@ -95,7 +95,7 @@ export function useDataStudioChatAgent() { acc[sessionSource.alias] = { connectionId: String(connection.id), dbType: attachedSource.databaseType, - permissions: sessionSource.permissions, + permissions: sessionSource.permissionsMode === 'custom' ? sessionSource.permissions : undefined, } return acc }, diff --git a/src/composables/useSidebarChatAgent.ts b/src/composables/useSidebarChatAgent.ts index 456bffaa..ca67aa91 100644 --- a/src/composables/useSidebarChatAgent.ts +++ b/src/composables/useSidebarChatAgent.ts @@ -67,7 +67,6 @@ export function useSidebarChatAgent() { [activeTab.name || 'active']: { connectionId: activeTab.connectionId, dbType: 'SQL', - permissions: { read: true, create: false, update: false, delete: false }, }, } } diff --git a/src/datasources/agentApi.ts b/src/datasources/agentApi.ts index 1ace0f29..74b5fac6 100644 --- a/src/datasources/agentApi.ts +++ b/src/datasources/agentApi.ts @@ -247,6 +247,7 @@ export const agentApi = { session_id: string tool_call_id: string tool_name: string + risk_level?: 'safe' | 'elevated' | 'destructive' arguments: unknown }) => void, ): Promise<() => void> => diff --git a/src/lang/enUS.ts b/src/lang/enUS.ts index 4497cc1c..65287d1d 100644 --- a/src/lang/enUS.ts +++ b/src/lang/enUS.ts @@ -1,4 +1,45 @@ export const enUS = { + aiPermissions: { + accessLevel: 'Access level', + unavailable: 'Unavailable', + loadError: 'Could not load permissions', + saving: 'Saving…', + conversationHelp: 'Uses the current connection permissions. Conversation limits can only restrict access; Ask/Auto only controls confirmation.', + useConnection: 'Use connection permissions', + limitConversation: 'Limit this conversation', + deniedByConnection: 'Denied by connection permissions', + cancel: 'Cancel', + conversationReadOnly: 'Read-only for this conversation', + title: 'AI permissions', + help: 'Applies to Data Studio and external MCP clients. Each explicit connection decision takes priority over the system default.', + loading: 'Loading permissions…', + loadFailed: 'Could not load permissions. Retry before saving this connection.', + retry: 'Retry', + inherit: 'Use system default', + none: 'No access', + read: 'Read', + data: 'Read & modify data', + full: 'Full access', + custom: 'Custom', + customize: 'Customize permissions', + effective: 'Effective access', + create: 'Create data', + update: 'Update data', + delete: 'Delete data', + manage: 'Manage structure & administration', + allow: 'Allow', + deny: 'Deny', + fromSystem: 'System default', + fromConnection: 'Connection setting', + defaults: 'Default AI permissions', + defaultsHelp: 'Used for actions that a connection inherits. Explicit connection decisions take priority. Conversation restrictions can only narrow access.', + save: 'Save defaults', + saved: 'Permissions saved', + saveFailed: 'Could not save permissions', + manageHelp: 'Full access also permits schema, index, table, collection, and administrative changes.', + readPreset: 'Read-only', + }, + device: { replaceTitle: 'Device limit reached', replaceDescription: @@ -1110,8 +1151,8 @@ export const enUS = { modifySource: { title: 'Modify Datasource', accessPermissions: 'Access Permissions', - inheritTitle: 'Inherit from session', - inheritDesc: 'Permissions follow the session-level setting and update automatically.', + inheritTitle: 'Use connection permissions', + inheritDesc: 'Uses the current connection permission chain.', customTitle: 'Custom', customDesc: 'Manually define which operations are allowed for this source.', read: 'Read', diff --git a/src/lang/zhCN.ts b/src/lang/zhCN.ts index 13ec79e6..5880b3d8 100644 --- a/src/lang/zhCN.ts +++ b/src/lang/zhCN.ts @@ -1,4 +1,45 @@ export const zhCN = { + aiPermissions: { + accessLevel: '访问级别', + unavailable: '暂不可用', + loadError: '无法加载权限', + saving: '正在保存…', + conversationHelp: '使用当前连接权限。会话设置只能进一步限制访问；询问/自动模式只控制操作确认。', + useConnection: '使用连接权限', + limitConversation: '限制当前会话', + deniedByConnection: '连接权限已禁止', + cancel: '取消', + conversationReadOnly: '当前会话只读', + title: 'AI 权限', + help: '适用于 Data Studio 和外部 MCP 客户端。连接中明确指定的每项权限优先于系统默认值。', + loading: '正在加载权限…', + loadFailed: '无法加载权限。请重试后再保存连接。', + retry: '重试', + inherit: '使用系统默认值', + none: '禁止访问', + read: '读取', + data: '读取和修改数据', + full: '完全访问', + custom: '自定义', + customize: '自定义权限', + effective: '实际权限', + create: '创建数据', + update: '更新数据', + delete: '删除数据', + manage: '管理结构和管理操作', + allow: '允许', + deny: '禁止', + fromSystem: '系统默认值', + fromConnection: '连接设置', + defaults: '默认 AI 权限', + defaultsHelp: '用于连接选择继承的操作权限。连接中明确指定的权限优先。会话权限只能进一步限制访问。', + save: '保存默认权限', + saved: '权限已保存', + saveFailed: '无法保存权限', + manageHelp: '完全访问还允许更改数据库结构、索引、表、集合及执行管理操作。', + readPreset: '只读', + }, + device: { replaceTitle: '设备数量已达上限', replaceDescription: diff --git a/src/pages/DataStudioPage.vue b/src/pages/DataStudioPage.vue index c5fdcf1e..8a15a614 100644 --- a/src/pages/DataStudioPage.vue +++ b/src/pages/DataStudioPage.vue @@ -451,7 +451,7 @@ function syncAllProviderModels() { @click="addSourceMode = 'Ask'" > - {{ t('dataStudio.modifySource.modeDefault') }} + {{ t('aiPermissions.conversationReadOnly') }}