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..9aa85622 --- /dev/null +++ b/docs/AI_PERMISSIONS.md @@ -0,0 +1,25 @@ +# 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 agent runtime to an exact Git revision in their Cargo manifests and lockfiles. Each application owns its permission policy and database-operation classification. + +The application-owned `ai_permissions` module resolves policy, migrates legacy settings, classifies database payloads, and gates both MCP and built-in agent execution. The shared agent routes requests through an application preparation hook and carries its confirmation risk and context; it does not resolve this permission chain. Execution rechecks the application policy on every attempt. + +AI schema reads use the same application gate as tool calls. Cached schema and editor context are included only when the associated connections permit reading. Manual database UI commands retain their existing access behavior. SQLKit does not expose a generic capability invocation command to the WebView. + +Agent executors are bound to a conversation and its connection/source identities. Before dispatch they reload connection policy, current attachment restrictions, and confirmation rules. Detaching or deleting a source prevents subsequent dispatches. The sidebar uses its bound active connection without requiring a Data Studio attachment. + +Preparation rejects denied calls before displaying confirmation. A queued call whose confirmation requirements become stricter returns to the existing approval flow. Permission denials are terminal and do not trigger transport retries; approval never bypasses a later denial. Rechecks occur at dispatch and do not cancel operations already dispatched or erase context already sent to a model. + +The host-hook API is introduced by [data-studio-agent #35](https://github.com/geek-fun/data-studio-agent/pull/35). Merge that runtime implementation before merging its consumers. Keep the exact compatible revision until a release containing these hooks exists; after the runtime is merged and released, update both the manifest and lockfile to that release. Older release tags do not provide this API. diff --git a/package-lock.json b/package-lock.json index 30f4f49f..e581d288 100644 --- a/package-lock.json +++ b/package-lock.json @@ -22,7 +22,7 @@ "@types/dagre": "^0.7.54", "@vueuse/core": "^14.1.0", "dagre": "^0.8.5", - "dompurify": "^3.4.10", + "dompurify": "^3.4.17", "highlight.js": "^11.11.1", "json-with-bigint": "^3.5.8", "json5": "^2.2.3", @@ -5879,9 +5879,9 @@ } }, "node_modules/dompurify": { - "version": "3.4.16", - "resolved": "https://registry.npmjs.org/dompurify/-/dompurify-3.4.16.tgz", - "integrity": "sha512-sqo+pNp3qRhCIpbgRi1y8Tgk27Bo2Ry7w0dC1NBeNTdZChWjz9Xb/KOoZbRP/R6pQZ80Qw8YhXw13hWWBbMRnQ==", + "version": "3.4.17", + "resolved": "https://registry.npmjs.org/dompurify/-/dompurify-3.4.17.tgz", + "integrity": "sha512-B/pSlYVHQ58kqISD/yDjuJ1yfzI34m1gF4lAod2bwSRhiN4K5Cl3H8Af7hqB2PmTL911tdEudnnXn5FnhRoz0w==", "license": "(MPL-2.0 OR Apache-2.0)", "optionalDependencies": { "@types/trusted-types": "^2.0.7" @@ -10621,15 +10621,6 @@ "marked": "14.0.0" } }, - "node_modules/monaco-editor/node_modules/dompurify": { - "version": "3.2.7", - "resolved": "https://registry.npmjs.org/dompurify/-/dompurify-3.2.7.tgz", - "integrity": "sha512-WhL/YuveyGXJaerVlMYGWhvQswa7myDG17P7Vu65EWC05o8vfeNbvNf4d/BOvH99+ZW+LlQsc1GDKMa1vNK6dw==", - "license": "(MPL-2.0 OR Apache-2.0)", - "optionalDependencies": { - "@types/trusted-types": "^2.0.7" - } - }, "node_modules/moo": { "version": "0.5.3", "resolved": "https://registry.npmjs.org/moo/-/moo-0.5.3.tgz", diff --git a/package.json b/package.json index dcf89398..1037aa75 100644 --- a/package.json +++ b/package.json @@ -82,7 +82,7 @@ "@types/dagre": "^0.7.54", "@vueuse/core": "^14.1.0", "dagre": "^0.8.5", - "dompurify": "^3.4.10", + "dompurify": "^3.4.17", "highlight.js": "^11.11.1", "json-with-bigint": "^3.5.8", "json5": "^2.2.3", @@ -132,5 +132,8 @@ "unocss-preset-shadcn": "^1.0.1", "vite": "^6.0.3", "vue-tsc": "^2.1.10" + }, + "overrides": { + "dompurify": "$dompurify" } } diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index 265c8880..cf7f4695 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=dc58625cdc721b64ef46a27bd6dcf1257fdd6f7b#dc58625cdc721b64ef46a27bd6dcf1257fdd6f7b" dependencies = [ "async-openai", "async-trait", diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index ca52287a..1f8a1aa3 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -87,8 +87,9 @@ 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 agent runtime to a reproducible revision. +# Requires the host hooks from data-studio-agent #35; older release tags are incompatible. +data-studio-agent = { git = "https://github.com/geek-fun/data-studio-agent.git", rev = "dc58625cdc721b64ef46a27bd6dcf1257fdd6f7b" } # 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..012d21c8 100644 --- a/src-tauri/src/agent/executor.rs +++ b/src-tauri/src/agent/executor.rs @@ -1,6 +1,9 @@ use async_trait::async_trait; -use data_studio_agent::tool_executor::{ToolEnvelope, ToolExecutor, ToolResultMetadata}; +use data_studio_agent::tool_executor::{ + ToolCallPreparation, ToolEnvelope, ToolExecutionError, ToolExecutor, ToolResultMetadata, +}; use serde_json::Value; +use std::collections::HashMap; const TOOL_ENVELOPE_MAX_CHARS: usize = 32768; const TOOL_ENVELOPE_SUMMARY_CHARS: usize = 4096; @@ -13,31 +16,47 @@ fn char_truncate(input: &str, max_chars: usize) -> (String, bool) { (truncated, true) } -pub struct SqlKitToolExecutor; +pub(crate) struct SqlKitToolExecutor { + pub(crate) access: crate::ai_permissions::AgentAccess, +} #[async_trait] impl ToolExecutor for SqlKitToolExecutor { + fn prepare( + &self, + tool_name: &str, + arguments: &Value, + config: &Value, + _connections: &HashMap, + ) -> Result { + crate::ai_permissions::prepare_agent_call(&self.access, tool_name, arguments, config) + } + async fn execute( &self, tool_name: &str, arguments: &Value, connection_config: &Value, ) -> Result { - let start = std::time::Instant::now(); - - let conn_opt = if connection_config.is_null() { - None - } else { - Some(connection_config.clone()) - }; + self.execute_checked(tool_name, arguments, connection_config) + .await + .map_err(|e| e.to_string()) + } - let raw = data_studio_agent::capabilities::registry::invoke_capability_inner( + async fn execute_checked( + &self, + tool_name: &str, + arguments: &Value, + connection_config: &Value, + ) -> Result { + let start = std::time::Instant::now(); + let raw = crate::ai_permissions::invoke_ai( tool_name, - arguments.clone(), - conn_opt, + arguments, + connection_config, + Some(&self.access), ) .await?; - 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..28034a36 100644 --- a/src-tauri/src/agent_adapters.rs +++ b/src-tauri/src/agent_adapters.rs @@ -5,9 +5,7 @@ //! to the shared library. 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}; @@ -49,9 +47,6 @@ pub async fn run_agent_loop( let confirm_map: ConfirmMap = confirm_state.inner().clone(); let cancel_state: State = app.state::(); let cancel_map: CancelMap = cancel_state.inner().clone(); - let executor_state: State> = - app.state::>(); - let executor: Arc = executor_state.inner().clone(); let connections: HashMap = settings .get("connections") @@ -62,6 +57,20 @@ pub async fn run_agent_loop( .cloned() .unwrap_or(Value::Null); + let studio = match settings.get("contextKind").and_then(Value::as_str) { + Some("dataStudio") => true, + Some("sidebarAssistant") => false, + _ => return Err("Unknown AI context".into()), + }; + let executor = crate::agent::executor::SqlKitToolExecutor { + access: crate::ai_permissions::AgentAccess { + session_id: session_id.clone(), + connections: connections.clone(), + studio, + db: db_state.inner().clone(), + }, + }; + let is_parallel_ok = |name: &str| -> bool { data_studio_agent::capabilities::registry::registry() .get(name) @@ -69,30 +78,8 @@ 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 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, - } - }; + // Required by the shared loop API; authorization is owned by the executor at both stages. + let legacy_deny_hook = |_: &str, _: Option<&str>| false; lib::loop_runner::run_agent_loop( &session_id, @@ -100,13 +87,13 @@ pub async fn run_agent_loop( &settings, &store, &emitter, - executor.as_ref(), + &executor, connections, fallback, &confirm_map, &cancel_map, &is_parallel_ok, - &should_deny, + &legacy_deny_hook, ) .await } diff --git a/src-tauri/src/ai_permissions/chain.rs b/src-tauri/src/ai_permissions/chain.rs new file mode 100644 index 00000000..5ca3fd8f --- /dev/null +++ b/src-tauri/src/ai_permissions/chain.rs @@ -0,0 +1,480 @@ +use super::{policy::PolicyAction, AiPolicy}; +use data_studio_agent::capabilities::types::{Capability, RiskLevel, SourceKind}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Permission { + Read, + Create, + Update, + Delete, + Manage, +} + +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct PermissionRules { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub read: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub create: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub update: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub delete: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub manage: Option, +} + +impl PermissionRules { + pub fn get(&self, action: Permission) -> Option { + match action { + Permission::Read => self.read, + Permission::Create => self.create, + Permission::Update => self.update, + Permission::Delete => self.delete, + Permission::Manage => self.manage, + } + } + + pub fn explicit(read: bool, create: bool, update: bool, delete: bool, manage: bool) -> Self { + Self { + read: Some(read), + create: Some(create), + update: Some(update), + delete: Some(delete), + manage: Some(manage), + } + } + + pub fn resolve(&self, defaults: &Self) -> Self { + Self { + read: Some(self.read.or(defaults.read).unwrap_or(false)), + create: Some(self.create.or(defaults.create).unwrap_or(false)), + update: Some(self.update.or(defaults.update).unwrap_or(false)), + delete: Some(self.delete.or(defaults.delete).unwrap_or(false)), + manage: Some(self.manage.or(defaults.manage).unwrap_or(false)), + } + } + + pub fn any_allowed(&self) -> bool { + [ + self.read, + self.create, + self.update, + self.delete, + self.manage, + ] + .contains(&Some(true)) + } +} + +impl AiPolicy { + pub fn advertises(&self, cap: &Capability) -> bool { + if self.permission_defaults.is_none() { + return self.allows(cap.risk_level, None); + } + if matches!(cap.source_kind, SourceKind::AppLocal | SourceKind::File) { + return true; + } + self.connection_permissions + .keys() + .map(String::as_str) + .chain(std::iter::once("__inherited__")) + .any(|id| { + let rules = self.effective_permissions(id); + if cap.name.ends_with("__execute_write") { + return rules.create == Some(true) || rules.update == Some(true); + } + required_permissions(cap, &serde_json::json!({})) + .is_ok_and(|actions| actions.iter().all(|a| rules.get(*a) == Some(true))) + }) + } + + pub fn effective_permissions(&self, id: &str) -> PermissionRules { + let Some(defaults) = &self.permission_defaults else { + let read = self.allows(RiskLevel::Safe, Some(id)); + let write = self.allows(RiskLevel::Elevated, Some(id)); + let delete = self.allows(RiskLevel::Destructive, Some(id)); + return PermissionRules::explicit(read, write, write, delete, delete); + }; + self.connection_permissions + .get(id) + .cloned() + .unwrap_or_default() + .resolve(defaults) + } + + /// Freeze legacy exceptions to their effective decisions before retiring the allowlist. + pub fn migrate_permissions(&mut self, connection_ids: &[String]) { + if self.permission_defaults.is_some() { + return; + } + let mut migrated = std::collections::HashMap::new(); + for id in connection_ids { + if !self.allowed_connection_ids.is_empty() || self.connection_overrides.contains_key(id) + { + migrated.insert(id.clone(), self.effective_permissions(id)); + } + } + let read = self.allows(RiskLevel::Safe, None); + let write = self.allows(RiskLevel::Elevated, None); + let delete = self.allows(RiskLevel::Destructive, None); + self.permission_defaults = Some(PermissionRules::explicit( + read, write, write, delete, delete, + )); + // A legacy allowlist also excluded future connections. + if !self.allowed_connection_ids.is_empty() { + self.permission_defaults = + Some(PermissionRules::explicit(false, false, false, false, false)); + } + self.connection_permissions = migrated; + self.allowed_connection_ids.clear(); + self.connection_overrides.clear(); + } + + pub fn authorize( + &self, + cap: &Capability, + args: &Value, + id: Option<&str>, + restriction: Option<&PermissionRules>, + ) -> Result { + if matches!(cap.source_kind, SourceKind::AppLocal) { + return Ok(PolicyAction::Allow); + } + if matches!(cap.source_kind, SourceKind::File) { + return if restriction.and_then(|r| r.read) == Some(false) { + Err("Read access to this source is denied".into()) + } else { + Ok(PolicyAction::Allow) + }; + } + let id = id.ok_or("A saved connection must be specified for database operations")?; + let effective = self.effective_permissions(id); + if !effective.any_allowed() { + return Err(format!("AI access denied for connection '{id}'")); + } + let actions = required_permissions(cap, args)?; + let requires_confirmation = matches!(cap.risk_level, RiskLevel::Destructive) + || actions.contains(&Permission::Delete) + || actions.contains(&Permission::Manage); + for action in actions { + if effective.get(action) != Some(true) { + return Err(format!( + "{action:?} access denied for connection '{id}' by AI permissions" + )); + } + if let Some(rules) = restriction { + if rules.get(action) == Some(false) + || (action == Permission::Manage && rules.manage != Some(true)) + { + return Err(format!( + "{action:?} access denied by conversation permissions" + )); + } + } + } + Ok(if requires_confirmation { + PolicyAction::Ask + } else { + PolicyAction::Allow + }) + } +} + +/// Payload-dependent tools require every action they can perform, rather than using risk as permission. +pub fn required_permissions(cap: &Capability, args: &Value) -> Result, String> { + let name = cap.name; + let op = name.rsplit("__").next().unwrap_or(name); + if op == "execute_write" { + use sqlparser::keywords::Keyword; + use sqlparser::tokenizer::{Token, Tokenizer}; + let sql = args + .get("sql") + .or_else(|| args.get("statement")) + .and_then(Value::as_str) + .ok_or("Missing SQL statement")?; + let tokens = Tokenizer::new(&sqlparser::dialect::GenericDialect {}, sql) + .tokenize() + .map_err(|e| e.to_string())?; + let mut actions = Vec::new(); + for token in tokens { + if let Token::Word(word) = token { + if word.quote_style.is_some() { + continue; + } + match word.keyword { + Keyword::INSERT => actions.push(Permission::Create), + Keyword::UPDATE => actions.push(Permission::Update), + Keyword::DELETE => actions.push(Permission::Delete), + Keyword::MERGE => { + actions.push(Permission::Create); + actions.push(Permission::Update); + actions.push(Permission::Delete); + } + _ => {} + } + } + } + if actions.is_empty() { + return Err("No supported SQL write operation".into()); + } + return Ok(actions); + } + if op == "execute_ddl" + || [ + "index", + "table", + "collection", + "database", + "mapping", + "alias", + "snapshot", + "repository", + "privilege", + "session", + "user", + "role", + ] + .iter() + .any(|object| { + (op.starts_with("create_") + || op.starts_with("delete_") + || op.starts_with("drop_") + || op.starts_with("put_") + || op.starts_with("update_") + || op.starts_with("restore_") + || op.starts_with("grant_") + || op.starts_with("revoke_") + || op.starts_with("kill_")) + && op.contains(object) + }) + { + return Ok(vec![Permission::Manage]); + } + match cap.required_permission { + "read" => Ok(vec![Permission::Read]), + "create" => Ok(vec![Permission::Create]), + "update" => Ok(vec![Permission::Update]), + "delete" => Ok(vec![Permission::Delete]), + "manage" => Ok(vec![Permission::Manage]), + "write" => Ok(vec![Permission::Create, Permission::Update]), + "none" if matches!(cap.source_kind, SourceKind::AppLocal) => Ok(vec![]), + _ => Err("Capability has no supported permission classification".into()), + } +} + +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn resolves_each_action_independently() { + let defaults = PermissionRules::explicit(true, true, true, false, false); + let rules = PermissionRules { + update: Some(false), + delete: Some(true), + ..Default::default() + }; + assert_eq!( + rules.resolve(&defaults), + PermissionRules::explicit(true, true, false, true, false) + ); + assert_eq!( + PermissionRules::default().resolve(&PermissionRules::default()), + PermissionRules::explicit(false, false, false, false, false) + ); + } + #[test] + fn migration_preserves_allowlist_exclusions_and_future_default() { + let mut policy = AiPolicy { + allowed_connection_ids: vec!["dev".into()], + ..Default::default() + }; + let before = policy.effective_permissions("dev"); + policy.migrate_permissions(&["dev".into(), "prod".into()]); + assert_eq!(policy.effective_permissions("dev"), before); + assert!(!policy.effective_permissions("prod").any_allowed()); + assert!(!policy.effective_permissions("new").any_allowed()); + assert!(policy.allowed_connection_ids.is_empty()); + let again = policy.clone(); + policy.migrate_permissions(&["dev".into()]); + assert_eq!(policy, again); + } +} + +#[cfg(test)] +mod authorization_tests { + use super::*; + use std::sync::Arc; + struct NeverExecute; + #[async_trait::async_trait] + impl data_studio_agent::capabilities::types::CapabilityHandler for NeverExecute { + async fn handle(&self, _: &Value, _: Option<&Value>) -> Result { + panic!("permission tests must not execute tools") + } + } + fn capability(name: &'static str, permission: &'static str, risk: RiskLevel) -> Capability { + Capability { + name, + required_permission: permission, + risk_level: risk, + source_kind: SourceKind::SqlDatabase, + input_schema: serde_json::json!({}), + description: "test", + handler: Arc::new(NeverExecute), + tags: &["agent"], + parallel_ok: false, + } + } + fn policy() -> AiPolicy { + AiPolicy { + permission_defaults: Some(PermissionRules::explicit(true, false, false, false, false)), + ..Default::default() + } + } + #[test] + fn connection_decisions_win_per_action_and_conversation_only_narrows() { + let mut policy = policy(); + policy.connection_permissions.insert( + "dev".into(), + PermissionRules { + update: Some(true), + read: Some(false), + ..Default::default() + }, + ); + let update = capability("test__update_item", "update", RiskLevel::Elevated); + let read = capability("test__query", "read", RiskLevel::Safe); + assert!(policy + .authorize(&update, &serde_json::json!({}), Some("dev"), None) + .is_ok()); + assert!(policy + .authorize(&read, &serde_json::json!({}), Some("dev"), None) + .is_err()); + let restriction = PermissionRules { + update: Some(false), + read: Some(true), + ..Default::default() + }; + assert!(policy + .authorize( + &update, + &serde_json::json!({}), + Some("dev"), + Some(&restriction) + ) + .is_err()); + assert!(policy + .authorize( + &read, + &serde_json::json!({}), + Some("dev"), + Some(&restriction) + ) + .is_err()); + assert!(policy + .authorize(&read, &serde_json::json!({}), Some("other"), None) + .is_ok()); + assert!(policy + .authorize(&read, &serde_json::json!({}), None, None) + .is_err()); + } + #[test] + fn explicit_allow_can_override_a_denied_system_default() { + let mut policy = policy(); + policy.permission_defaults = + Some(PermissionRules::explicit(false, false, false, false, false)); + policy.connection_permissions.insert( + "dev".into(), + PermissionRules { + read: Some(true), + ..Default::default() + }, + ); + let read = capability("test__query", "read", RiskLevel::Safe); + assert!(policy.advertises(&read)); + assert!(policy + .authorize(&read, &serde_json::json!({}), Some("dev"), None) + .is_ok()); + assert!(!policy.effective_permissions("other").any_allowed()); + } + #[test] + fn data_delete_does_not_grant_schema_management() { + let mut policy = policy(); + policy.permission_defaults = Some(PermissionRules::explicit(true, true, true, true, false)); + assert!(policy + .authorize( + &capability("sqlkit__execute_delete", "delete", RiskLevel::Destructive), + &serde_json::json!({}), + Some("dev"), + None + ) + .is_ok()); + assert!(policy + .authorize( + &capability("sqlkit__execute_ddl", "delete", RiskLevel::Destructive), + &serde_json::json!({}), + Some("dev"), + None + ) + .is_err()); + } + #[test] + fn persistence_preserves_explicit_false_and_inherited_omission() { + let rules = PermissionRules { + update: Some(false), + ..Default::default() + }; + assert_eq!( + serde_json::to_value(&rules).unwrap(), + serde_json::json!({"update":false}) + ); + assert_eq!( + serde_json::from_value::(serde_json::json!({"update":false})).unwrap(), + rules + ); + assert!( + serde_json::from_value::(serde_json::json!({"execute":true})).is_err() + ); + } +} + +#[cfg(test)] +mod mixed_operation_tests { + use super::*; + use std::sync::Arc; + struct Noop; + #[async_trait::async_trait] + impl data_studio_agent::capabilities::types::CapabilityHandler for Noop { + async fn handle(&self, _: &Value, _: Option<&Value>) -> Result { + Ok(String::new()) + } + } + fn cap(name: &'static str, permission: &'static str) -> Capability { + Capability { + name, + required_permission: permission, + risk_level: RiskLevel::Elevated, + source_kind: SourceKind::SqlDatabase, + input_schema: serde_json::json!({}), + description: "", + handler: Arc::new(Noop), + tags: &["agent"], + parallel_ok: false, + } + } + #[test] + fn sql_ignores_comments_and_string_literals() { + assert_eq!( + required_permissions( + &cap("sqlkit__execute_write", "create"), + &serde_json::json!({"sql":"-- UPDATE ignored\nINSERT INTO t VALUES ('DELETE');"}) + ) + .unwrap(), + vec![Permission::Create] + ); + } +} diff --git a/src-tauri/src/ai_permissions/invocation.rs b/src-tauri/src/ai_permissions/invocation.rs new file mode 100644 index 00000000..7d39e864 --- /dev/null +++ b/src-tauri/src/ai_permissions/invocation.rs @@ -0,0 +1,376 @@ +use super::{load_permission_policy, policy::PolicyAction, scope::AgentAccess, AiPolicy}; +use data_studio_agent::{ + capabilities::{ + registry, + types::{Capability, RiskLevel}, + }, + tool_executor::{ToolCallPreparation, ToolExecutionError}, +}; +use serde_json::{json, Value}; +use tauri::Manager; +use tauri_plugin_store::StoreExt; + +struct Authorization { + policy: AiPolicy, + risk: RiskLevel, + confirmation: bool, + allowed_ids: Option>, +} + +fn authorize( + name: &str, + args: &Value, + config: &Value, + access: Option<&AgentAccess>, +) -> Result { + let app = crate::APP_HANDLE.get().ok_or("App not initialized")?; + if let Some(id) = config.get("connectionId").and_then(Value::as_str) { + let connections = app + .store(".store.dat") + .map_err(|e| e.to_string())? + .get("connections") + .and_then(|value| value.as_array().cloned()) + .unwrap_or_default(); + if !connections + .iter() + .any(|connection| match &connection["id"] { + Value::String(value) => value == id, + Value::Number(value) => value.to_string() == id, + _ => false, + }) + { + return Err("Saved connection is no longer available".into()); + } + } + let policy = load_permission_policy(app)?; + let cap = registry::registry().get(name).ok_or("Unknown capability")?; + let session = access.map(|access| access.current(name)).transpose()?; + evaluate(policy, cap, args, config, session) +} + +fn evaluate( + policy: AiPolicy, + cap: &Capability, + args: &Value, + config: &Value, + session: Option, +) -> Result { + let id = config.get("connectionId").and_then(Value::as_str); + let restriction = match (id, &session) { + (Some(id), Some(session)) => Some( + session + .restrictions + .get(id) + .ok_or("Connection is no longer attached to this conversation")?, + ), + _ => None, + }; + let decision = policy.authorize(cap, args, id, restriction)?; + let risk = if decision == PolicyAction::Ask { + RiskLevel::Destructive + } else { + cap.risk_level + }; + let confirmation = session + .as_ref() + .map(|session| session.confirmation_required(risk)) + .transpose()? + .unwrap_or(false); + let allowed_ids = session.map(|session| { + session + .restrictions + .into_iter() + .filter_map(|(id, restriction)| { + let effective = policy.effective_permissions(&id); + [ + super::chain::Permission::Read, + super::chain::Permission::Create, + super::chain::Permission::Update, + super::chain::Permission::Delete, + super::chain::Permission::Manage, + ] + .iter() + .any(|action| { + effective.get(*action) == Some(true) + && restriction.get(*action) != Some(false) + && (*action != super::chain::Permission::Manage + || restriction.manage == Some(true)) + }) + .then_some(id) + }) + .collect() + }); + Ok(Authorization { + policy, + risk, + confirmation, + allowed_ids, + }) +} + +pub(crate) fn prepare_agent_call( + access: &AgentAccess, + name: &str, + args: &Value, + config: &Value, +) -> Result { + let authorization = authorize(name, args, config, Some(access))?; + let mut config = if config.is_object() { + config.clone() + } else { + json!({}) + }; + config["__ai_confirmation_required"] = json!(authorization.confirmation); + Ok(ToolCallPreparation { + connection_config: config, + risk_level: Some(authorization.risk), + confirmation_required: Some(authorization.confirmation), + }) +} + +fn check_confirmation(config: &Value, required: bool) -> Result<(), ToolExecutionError> { + if required + && config + .get("__ai_confirmation_required") + .and_then(Value::as_bool) + != Some(true) + { + return Err(ToolExecutionError::ConfirmationRequired); + } + Ok(()) +} + +pub(crate) async fn invoke_ai( + name: &str, + args: &Value, + config: &Value, + access: Option<&AgentAccess>, +) -> Result { + let initial = authorize(name, args, config, access).map_err(ToolExecutionError::Denied)?; + if access.is_some() { + check_confirmation(config, initial.confirmation)?; + } + let resolved = match config.get("connectionId").and_then(Value::as_str) { + Some(id) => Some(resolve_connection(id).await?), + None => None, + }; + // Resolution can await connection/SSH setup. Revalidate at dispatch after it completes. + let authorization = + authorize(name, args, config, access).map_err(ToolExecutionError::Denied)?; + if access.is_some() { + check_confirmation(config, authorization.confirmation)?; + } + let raw = registry::invoke_capability_inner(name, args.clone(), resolved).await?; + if !name.ends_with("__list_connections") { + return Ok(raw); + } + 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) { + result["connections"] = Value::Array( + connections + .iter() + .filter(|connection| { + let id = match &connection["id"] { + Value::String(id) => id.clone(), + Value::Number(id) => id.to_string(), + _ => return false, + }; + authorization + .policy + .effective_permissions(&id) + .any_allowed() + && authorization + .allowed_ids + .as_ref() + .is_none_or(|ids| ids.contains(&id)) + }) + .cloned() + .collect(), + ); + } + Ok(result.to_string()) +} + +async fn resolve_connection(id: &str) -> Result { + crate::capabilities::sql::resolve_adapter(id).await?; + Ok(json!({"connectionId": id})) +} + +#[tauri::command] +pub async fn get_ai_context( + session_id: String, + connections: std::collections::HashMap, + context_kind: String, + app: tauri::AppHandle, +) -> Result { + crate::entitlement::ensure_local_ultimate( + &app.state::(), + "AI", + )?; + let studio = match context_kind.as_str() { + "dataStudio" => true, + "sidebarAssistant" => false, + _ => return Err("Unknown AI context".into()), + }; + let access = AgentAccess { + session_id, + connections, + studio, + db: app + .state::() + .inner() + .clone(), + }; + let mut readable = Vec::new(); + let mut schemas = Vec::new(); + for (alias, config) in &access.connections { + let Some(id) = config.get("connectionId").and_then(Value::as_str) else { + continue; + }; + if readable.iter().any(|value| value == id) { + continue; + } + let schema = invoke_ai("sqlkit__get_schema", &json!({}), config, Some(&access)).await; + if matches!( + schema, + Err(ToolExecutionError::Denied(_) | ToolExecutionError::ConfirmationRequired) + ) { + continue; + } + // Refresh access after the read before including schema or cached editor context. + if authorize("sqlkit__get_schema", &json!({}), config, Some(&access)).is_err() { + continue; + } + readable.push(id.to_owned()); + if let Ok(schema) = schema { + schemas.push(format!("-- Source: {alias}\n{schema}")); + } + } + Ok(json!({"readableConnectionIds": readable, "schema": schemas.join("\n")})) +} + +#[cfg(test)] +mod tests { + use super::super::{scope::SessionAccess, PermissionRules}; + use super::*; + use data_studio_agent::capabilities::types::{Capability, CapabilityHandler, SourceKind}; + use std::{collections::HashMap, sync::Arc}; + + struct NeverExecute; + #[async_trait::async_trait] + impl CapabilityHandler for NeverExecute { + async fn handle(&self, _: &Value, _: Option<&Value>) -> Result { + panic!("authorization must not execute a capability") + } + } + fn cap() -> Capability { + Capability { + name: "sqlkit__execute_write", + required_permission: "create", + source_kind: SourceKind::SqlDatabase, + risk_level: RiskLevel::Elevated, + handler: Arc::new(NeverExecute), + description: "", + input_schema: json!({}), + tags: &["agent"], + parallel_ok: false, + } + } + fn policy() -> AiPolicy { + AiPolicy { + permission_defaults: Some(PermissionRules::explicit(true, true, true, true, true)), + ..Default::default() + } + } + fn session() -> SessionAccess { + SessionAccess { + restrictions: HashMap::from([( + "42".into(), + PermissionRules::explicit(true, true, true, true, true), + )]), + mode: "Auto".into(), + rule: None, + } + } + #[test] + fn connection_revocation_denies_even_after_successful_preflight() { + let cap = cap(); + let config = json!({"connectionId":"42"}); + let args = json!({"sql":"UPDATE t SET x=1"}); + let mut policy = policy(); + assert!(evaluate(policy.clone(), &cap, &args, &config, Some(session())).is_ok()); + policy.connection_permissions.insert( + "42".into(), + PermissionRules { + update: Some(false), + ..Default::default() + }, + ); + assert!(evaluate(policy, &cap, &args, &config, Some(session())).is_err()); + } + #[test] + fn confirmation_never_bypasses_connection_or_conversation_denial() { + let config = json!({"connectionId":"42","__ai_confirmation_required":true}); + let mut session = session(); + session.restrictions.get_mut("42").unwrap().update = Some(false); + assert!(evaluate( + policy(), + &cap(), + &json!({"sql":"UPDATE t SET x=1"}), + &config, + Some(session) + ) + .is_err()); + } + #[test] + fn unbound_targets_are_denied_but_mcp_has_no_conversation_attachment_requirement() { + let config = json!({"connectionId":"99"}); + let args = json!({"sql":"UPDATE t SET x=1"}); + assert!(evaluate(policy(), &cap(), &args, &config, Some(session())).is_err()); + assert!(evaluate(policy(), &cap(), &args, &config, None).is_ok()); + } + #[test] + fn payload_controls_confirmation_and_all_required_actions() { + let args = json!({"sql":"INSERT INTO t VALUES(1); DELETE FROM t"}); + let config = json!({"connectionId":"42"}); + let decision = evaluate(policy(), &cap(), &args, &config, Some(session())).unwrap(); + assert_eq!(decision.risk, RiskLevel::Destructive); + assert!(decision.confirmation); + let mut policy = policy(); + policy.connection_permissions.insert( + "42".into(), + PermissionRules { + delete: Some(false), + ..Default::default() + }, + ); + assert!(evaluate(policy, &cap(), &args, &config, Some(session())).is_err()); + } + #[test] + fn stricter_confirmation_requires_renewed_approval() { + assert!(check_confirmation(&json!({"__ai_confirmation_required":false}), false).is_ok()); + assert_eq!( + check_confirmation(&json!({"__ai_confirmation_required":false}), true), + Err(ToolExecutionError::ConfirmationRequired) + ); + assert!(check_confirmation(&json!({"__ai_confirmation_required":true}), true).is_ok()); + } + #[test] + fn discovery_is_limited_to_current_scope_and_effective_actions() { + let mut cap = cap(); + cap.name = "app__list_connections"; + cap.source_kind = SourceKind::AppLocal; + let mut session = session(); + session.restrictions.insert( + "42".into(), + PermissionRules::explicit(false, false, false, false, false), + ); + assert!( + evaluate(policy(), &cap, &json!({}), &Value::Null, Some(session)) + .unwrap() + .allowed_ids + .unwrap() + .is_empty() + ); + } +} diff --git a/src-tauri/src/ai_permissions/mod.rs b/src-tauri/src/ai_permissions/mod.rs new file mode 100644 index 00000000..51898c8c --- /dev/null +++ b/src-tauri/src/ai_permissions/mod.rs @@ -0,0 +1,95 @@ +mod chain; +pub(crate) mod invocation; +mod policy; +mod scope; +pub(crate) use invocation::{invoke_ai, prepare_agent_call}; +pub(crate) use scope::AgentAccess; + +pub(crate) use chain::PermissionRules; +pub(crate) use policy::AiPolicy; +#[cfg(test)] +pub(crate) use policy::McpPermissionMode; + +use crate::mcp_bridge::McpConfig; +use serde_json::Value; +use std::path::Path; +use tauri::{AppHandle, Manager}; +use tauri_plugin_store::StoreExt; + +pub(crate) 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) +} + +pub(crate) 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) +} diff --git a/src-tauri/src/ai_permissions/policy.rs b/src-tauri/src/ai_permissions/policy.rs new file mode 100644 index 00000000..987dc083 --- /dev/null +++ b/src-tauri/src/ai_permissions/policy.rs @@ -0,0 +1,790 @@ +use std::collections::HashMap; + +use serde::{Deserialize, Serialize}; + +use super::chain::PermissionRules; +use data_studio_agent::capabilities::types::RiskLevel; + +/// Result of a policy decision for a capability invocation. +/// +/// `Ask` is informational this iteration: both the MCP bridge and the built-in +/// agent loop execute on `Ask` (confirmation UX is client-driven). Only `Deny` +/// short-circuits execution. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum PolicyAction { + Allow, + Ask, + Deny, +} + +/// Permission modes for MCP bridge access. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum McpPermissionMode { + /// Read-only queries and metadata — only Safe capabilities + ReadOnly, + /// Read + write (insert, update, create) — Safe + Elevated + DataReadWrite, + /// Everything including destructive ops — Safe + Elevated + Destructive + FullAccess, +} + +impl McpPermissionMode { + pub fn allows(&self, risk_level: RiskLevel) -> bool { + match self { + McpPermissionMode::ReadOnly => matches!(risk_level, RiskLevel::Safe), + McpPermissionMode::DataReadWrite => { + matches!(risk_level, RiskLevel::Safe | RiskLevel::Elevated) + } + McpPermissionMode::FullAccess => true, + } + } +} + +/// Per-connection override that takes precedence over the global mode. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct ConnectionMcpOverride { + /// Read-only for MCP: blocks Elevated/Destructive regardless of global mode + pub read_only: bool, + /// Optional action-level allowlist for this connection. + /// When set, only capabilities whose risk maps to a listed action pass; + /// `read_only` remains the coarse legacy gate. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub allowed_actions: Option>, +} + +/// Action categories a connection-level override can allow. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum McpAction { + /// Safe capabilities (queries, metadata) + Read, + /// Elevated capabilities (insert, update, create) + Write, + /// Destructive capabilities (delete, drop, truncate) + Delete, +} + +impl McpAction { + pub fn from_risk(risk: RiskLevel) -> McpAction { + match risk { + RiskLevel::Safe => McpAction::Read, + RiskLevel::Elevated => McpAction::Write, + RiskLevel::Destructive => McpAction::Delete, + } + } +} + +/// MCP permission policy, stored per app and read by the bridge on every request. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct AiPolicy { + #[serde(default)] + pub permission_defaults: Option, + #[serde(default)] + pub connection_permissions: HashMap, + pub mode: McpPermissionMode, + /// Empty = all connections allowed; non-empty = only listed connection IDs + pub allowed_connection_ids: Vec, + /// Connection-ID-keyed overrides that supersede the global mode + pub connection_overrides: HashMap, + /// When true, Destructive capabilities are exposed and gated by mode + /// (only FullAccess allows them). When false, Destructive is always + /// hidden/rejected regardless of mode. + #[serde(default = "default_confirm_destructive")] + pub confirm_destructive: bool, +} + +fn default_confirm_destructive() -> bool { + true +} + +impl Default for AiPolicy { + fn default() -> Self { + Self { + permission_defaults: None, + connection_permissions: HashMap::new(), + mode: McpPermissionMode::DataReadWrite, + allowed_connection_ids: Vec::new(), + connection_overrides: HashMap::new(), + confirm_destructive: true, + } + } +} + +impl AiPolicy { + pub fn is_connection_allowed(&self, connection_id: &str) -> bool { + if self.permission_defaults.is_some() { + return self.effective_permissions(connection_id).any_allowed(); + } + self.allowed_connection_ids.is_empty() + || self + .allowed_connection_ids + .iter() + .any(|id| id == connection_id) + } + + pub fn is_connection_read_only(&self, connection_id: &str) -> bool { + if self.permission_defaults.is_some() { + let rules = self.effective_permissions(connection_id); + return rules.read == Some(true) + && [rules.create, rules.update, rules.delete, rules.manage] + .iter() + .all(|v| *v != Some(true)); + } + self.connection_overrides + .get(connection_id) + .is_some_and(|o| o.read_only) + } + + /// Whether the connection override allows the risk level's action. + /// Falls back to `read_only` when no action allowlist is set. + fn connection_allows_action(&self, connection_id: &str, risk: RiskLevel) -> bool { + let Some(ov) = self.connection_overrides.get(connection_id) else { + return true; + }; + match &ov.allowed_actions { + Some(actions) => actions.contains(&McpAction::from_risk(risk)), + None => !(ov.read_only && !matches!(risk, RiskLevel::Safe)), + } + } + + /// Whether a capability may run for the given connection. + /// Compat: `allows()` == `decide() != Deny` (Ask counts as allowed). + pub fn allows(&self, risk_level: RiskLevel, connection_id: Option<&str>) -> bool { + !matches!(self.decide(risk_level, connection_id), PolicyAction::Deny) + } + + /// Decide the policy action for a capability and connection. + pub fn decide(&self, risk_level: RiskLevel, connection_id: Option<&str>) -> PolicyAction { + if self.permission_defaults.is_some() { + let rules = self.effective_permissions(connection_id.unwrap_or("__inherited__")); + let allowed = match risk_level { + RiskLevel::Safe => rules.read == Some(true), + RiskLevel::Elevated => { + [rules.create, rules.update, rules.manage].contains(&Some(true)) + } + RiskLevel::Destructive => [rules.delete, rules.manage].contains(&Some(true)), + }; + return if !allowed { + PolicyAction::Deny + } else if matches!(risk_level, RiskLevel::Destructive) { + PolicyAction::Ask + } else { + PolicyAction::Allow + }; + } + if let Some(id) = connection_id { + if !self.is_connection_allowed(id) { + return PolicyAction::Deny; + } + if !self.connection_allows_action(id, risk_level) { + return PolicyAction::Deny; + } + } + if !self.confirm_destructive && matches!(risk_level, RiskLevel::Destructive) { + return PolicyAction::Deny; + } + if !self.mode.allows(risk_level) { + return PolicyAction::Deny; + } + if matches!(risk_level, RiskLevel::Destructive) && self.confirm_destructive { + PolicyAction::Ask + } else { + PolicyAction::Allow + } + } + + /// Human-readable, actionable reason a capability was denied. + /// Mirrors the `decide()` check order: connection allowlist → override + /// action → confirm_destructive → mode. Returns None when allowed (or Ask). + pub fn deny_reason( + &self, + risk_level: RiskLevel, + connection_id: Option<&str>, + ) -> Option { + if self.permission_defaults.is_some() { + return if self.allows(risk_level, connection_id) { + None + } else { + Some("Access denied by the resolved AI permission chain".into()) + }; + } + if let Some(id) = connection_id { + if !self.is_connection_allowed(id) { + return Some(format!( + "connection '{id}' is not in the MCP allowlist — add it in Settings → MCP Bridge" + )); + } + if !self.connection_allows_action(id, risk_level) { + return Some(format!( + "connection '{id}' is read-only or lacks the required action — adjust its override in Settings → MCP Bridge" + )); + } + } + if !self.confirm_destructive && matches!(risk_level, RiskLevel::Destructive) { + return Some( + "destructive operations are disabled — enable 'Confirm Destructive Operations' in Settings → MCP Bridge" + .into(), + ); + } + if !self.mode.allows(risk_level) { + let required = match risk_level { + RiskLevel::Safe => McpPermissionMode::ReadOnly, + RiskLevel::Elevated => McpPermissionMode::DataReadWrite, + RiskLevel::Destructive => McpPermissionMode::FullAccess, + }; + return Some(format!( + "requires {required:?} permission mode — switch it in Settings → MCP Bridge" + )); + } + None + } + + /// Policy notice appended to every tool description in `tools/list`, so + /// clients (LLMs) learn what is gated and how to lift the gate without + /// calling hidden tools (mirrors Neon MCP's read-only notice). + pub fn policy_notice(&self) -> Option { + if self.permission_defaults.is_some() { + return Some("AI permissions are resolved per action: explicit connection decisions take priority over system defaults. Conversation restrictions can only narrow access.".into()); + } + let mut parts: Vec = Vec::new(); + if !self.mode.allows(RiskLevel::Destructive) || !self.confirm_destructive { + parts.push( + "destructive operations (delete/drop/truncate) are hidden — enable Full Access with Confirm Destructive in Settings → MCP Bridge to use them" + .into(), + ); + } + if !self.mode.allows(RiskLevel::Elevated) { + parts.push( + "write operations (insert/update/create) are hidden — enable Data Read/Write or Full Access in Settings → MCP Bridge to use them" + .into(), + ); + } + if parts.is_empty() { + None + } else { + Some(format!("MCP policy notice: {}", parts.join(" "))) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn test_mode_allows_risk_levels() { + assert!(McpPermissionMode::ReadOnly.allows(RiskLevel::Safe)); + assert!(!McpPermissionMode::ReadOnly.allows(RiskLevel::Elevated)); + assert!(!McpPermissionMode::ReadOnly.allows(RiskLevel::Destructive)); + + assert!(McpPermissionMode::DataReadWrite.allows(RiskLevel::Safe)); + assert!(McpPermissionMode::DataReadWrite.allows(RiskLevel::Elevated)); + assert!(!McpPermissionMode::DataReadWrite.allows(RiskLevel::Destructive)); + + assert!(McpPermissionMode::FullAccess.allows(RiskLevel::Safe)); + assert!(McpPermissionMode::FullAccess.allows(RiskLevel::Elevated)); + assert!(McpPermissionMode::FullAccess.allows(RiskLevel::Destructive)); + } + + #[test] + fn test_default_policy_is_data_read_write_with_all_connections() { + let policy = AiPolicy::default(); + assert_eq!(policy.mode, McpPermissionMode::DataReadWrite); + assert!(policy.allowed_connection_ids.is_empty()); + assert!(policy.is_connection_allowed("any-id")); + assert!(!policy.is_connection_read_only("any-id")); + // Destructive confirmations are on by default (issue #10) + assert!(policy.confirm_destructive); + } + + #[test] + fn test_confirm_destructive_false_blocks_destructive_even_in_full_access() { + let policy = AiPolicy { + mode: McpPermissionMode::FullAccess, + confirm_destructive: false, + ..AiPolicy::default() + }; + // confirm_destructive=false hides Destructive from FullAccess too + assert!(!policy.allows(RiskLevel::Destructive, None)); + // Safe/Elevated unaffected + assert!(policy.allows(RiskLevel::Safe, None)); + assert!(policy.allows(RiskLevel::Elevated, None)); + } + + #[test] + fn test_confirm_destructive_true_allows_destructive_only_in_full_access() { + let policy = AiPolicy { + mode: McpPermissionMode::FullAccess, + confirm_destructive: true, + ..AiPolicy::default() + }; + assert!(policy.allows(RiskLevel::Destructive, None)); + + let readonly_policy = AiPolicy { + mode: McpPermissionMode::DataReadWrite, + confirm_destructive: true, + ..AiPolicy::default() + }; + // Mode gate still applies — DataReadWrite blocks Destructive regardless + assert!(!readonly_policy.allows(RiskLevel::Destructive, None)); + } + + #[test] + fn test_deny_reason_guides_mode_upgrade() { + let policy = AiPolicy::default(); + assert_eq!( + policy.deny_reason(RiskLevel::Destructive, None), + Some("requires FullAccess permission mode — switch it in Settings → MCP Bridge".into()) + ); + let read_only = AiPolicy { + mode: McpPermissionMode::ReadOnly, + ..AiPolicy::default() + }; + assert_eq!( + read_only.deny_reason(RiskLevel::Elevated, None), + Some( + "requires DataReadWrite permission mode — switch it in Settings → MCP Bridge" + .into() + ) + ); + // Allowed capabilities have no reason + assert_eq!(policy.deny_reason(RiskLevel::Safe, None), None); + assert_eq!(policy.deny_reason(RiskLevel::Elevated, None), None); + } + + #[test] + fn test_deny_reason_guides_confirm_destructive() { + let policy = AiPolicy { + mode: McpPermissionMode::FullAccess, + confirm_destructive: false, + ..AiPolicy::default() + }; + assert_eq!( + policy.deny_reason(RiskLevel::Destructive, None), + Some("destructive operations are disabled — enable 'Confirm Destructive Operations' in Settings → MCP Bridge".into()) + ); + } + + #[test] + fn test_deny_reason_guides_connection_override() { + let policy = AiPolicy { + mode: McpPermissionMode::FullAccess, + confirm_destructive: true, + allowed_connection_ids: vec!["conn-1".into()], + connection_overrides: [( + "conn-1".into(), + ConnectionMcpOverride { + read_only: true, + allowed_actions: None, + }, + )] + .into(), + ..AiPolicy::default() + }; + assert_eq!( + policy.deny_reason(RiskLevel::Safe, Some("conn-2")), + Some( + "connection 'conn-2' is not in the MCP allowlist — add it in Settings → MCP Bridge" + .into() + ) + ); + assert_eq!( + policy.deny_reason(RiskLevel::Elevated, Some("conn-1")), + Some("connection 'conn-1' is read-only or lacks the required action — adjust its override in Settings → MCP Bridge".into()) + ); + // FullAccess + confirm on + safe connection: no reason + assert_eq!(policy.deny_reason(RiskLevel::Safe, Some("conn-1")), None); + } + + #[test] + fn test_policy_notice_reflects_gated_risks() { + // Default DataReadWrite hides Destructive → notice mentions it + let default = AiPolicy::default(); + let notice = default.policy_notice().unwrap(); + assert!(notice.contains("destructive operations")); + assert!(!notice.contains("write operations")); + // ReadOnly hides both + let read_only = AiPolicy { + mode: McpPermissionMode::ReadOnly, + ..AiPolicy::default() + }; + let notice = read_only.policy_notice().unwrap(); + assert!(notice.contains("destructive operations")); + assert!(notice.contains("write operations")); + // FullAccess + confirm on hides nothing + let full = AiPolicy { + mode: McpPermissionMode::FullAccess, + confirm_destructive: true, + ..AiPolicy::default() + }; + assert_eq!(full.policy_notice(), None); + } + + #[test] + fn test_allowlist_restricts_connections() { + let policy = AiPolicy { + allowed_connection_ids: vec!["conn-1".into()], + ..AiPolicy::default() + }; + assert!(policy.is_connection_allowed("conn-1")); + assert!(!policy.is_connection_allowed("conn-2")); + assert!(!policy.allows(RiskLevel::Safe, Some("conn-2"))); + } + + #[test] + fn test_read_only_override_blocks_elevated_but_allows_safe() { + let policy = AiPolicy { + mode: McpPermissionMode::FullAccess, + connection_overrides: HashMap::from([( + "prod-es".into(), + ConnectionMcpOverride { + read_only: true, + allowed_actions: None, + }, + )]), + ..AiPolicy::default() + }; + assert!(policy.allows(RiskLevel::Destructive, Some("staging-es"))); + assert!(!policy.allows(RiskLevel::Destructive, Some("prod-es"))); + assert!(!policy.allows(RiskLevel::Elevated, Some("prod-es"))); + assert!(policy.allows(RiskLevel::Safe, Some("prod-es"))); + } + + #[test] + fn test_serde_roundtrip() { + let policy = AiPolicy { + mode: McpPermissionMode::DataReadWrite, + allowed_connection_ids: vec!["conn-1".into()], + connection_overrides: HashMap::from([( + "conn-1".into(), + ConnectionMcpOverride { + read_only: true, + allowed_actions: None, + }, + )]), + confirm_destructive: false, + ..AiPolicy::default() + }; + let v = serde_json::to_value(&policy).unwrap(); + assert_eq!(v["mode"], "DataReadWrite"); + assert_eq!(v["allowed_connection_ids"][0], "conn-1"); + assert_eq!(v["connection_overrides"]["conn-1"]["read_only"], true); + assert_eq!(v["confirm_destructive"], false); + + let back: AiPolicy = serde_json::from_value(v).unwrap(); + assert_eq!(back, policy); + } + + #[test] + fn test_mode_serde_names() { + assert_eq!( + serde_json::to_value(McpPermissionMode::ReadOnly).unwrap(), + json!("ReadOnly") + ); + assert_eq!( + serde_json::to_value(McpPermissionMode::DataReadWrite).unwrap(), + json!("DataReadWrite") + ); + assert_eq!( + serde_json::to_value(McpPermissionMode::FullAccess).unwrap(), + json!("FullAccess") + ); + } + + #[test] + fn test_decide_full_matrix() { + // mode × risk × confirm_destructive → expected PolicyAction + let cases: Vec<(McpPermissionMode, RiskLevel, bool, PolicyAction)> = vec![ + ( + McpPermissionMode::ReadOnly, + RiskLevel::Safe, + true, + PolicyAction::Allow, + ), + ( + McpPermissionMode::ReadOnly, + RiskLevel::Safe, + false, + PolicyAction::Allow, + ), + ( + McpPermissionMode::ReadOnly, + RiskLevel::Elevated, + true, + PolicyAction::Deny, + ), + ( + McpPermissionMode::ReadOnly, + RiskLevel::Elevated, + false, + PolicyAction::Deny, + ), + ( + McpPermissionMode::ReadOnly, + RiskLevel::Destructive, + true, + PolicyAction::Deny, + ), + ( + McpPermissionMode::ReadOnly, + RiskLevel::Destructive, + false, + PolicyAction::Deny, + ), + ( + McpPermissionMode::DataReadWrite, + RiskLevel::Safe, + true, + PolicyAction::Allow, + ), + ( + McpPermissionMode::DataReadWrite, + RiskLevel::Safe, + false, + PolicyAction::Allow, + ), + ( + McpPermissionMode::DataReadWrite, + RiskLevel::Elevated, + true, + PolicyAction::Allow, + ), + ( + McpPermissionMode::DataReadWrite, + RiskLevel::Elevated, + false, + PolicyAction::Allow, + ), + ( + McpPermissionMode::DataReadWrite, + RiskLevel::Destructive, + true, + PolicyAction::Deny, + ), + ( + McpPermissionMode::DataReadWrite, + RiskLevel::Destructive, + false, + PolicyAction::Deny, + ), + ( + McpPermissionMode::FullAccess, + RiskLevel::Safe, + true, + PolicyAction::Allow, + ), + ( + McpPermissionMode::FullAccess, + RiskLevel::Safe, + false, + PolicyAction::Allow, + ), + ( + McpPermissionMode::FullAccess, + RiskLevel::Elevated, + true, + PolicyAction::Allow, + ), + ( + McpPermissionMode::FullAccess, + RiskLevel::Elevated, + false, + PolicyAction::Allow, + ), + // NEW semantic: confirm_destructive=true + Destructive + FullAccess → Ask + ( + McpPermissionMode::FullAccess, + RiskLevel::Destructive, + true, + PolicyAction::Ask, + ), + ( + McpPermissionMode::FullAccess, + RiskLevel::Destructive, + false, + PolicyAction::Deny, + ), + ]; + + for (mode, risk, confirm, expected) in cases { + let policy = AiPolicy { + mode, + confirm_destructive: confirm, + ..AiPolicy::default() + }; + let decided = policy.decide(risk, None); + assert_eq!( + decided, expected, + "decide({mode:?}, {risk:?}, confirm_destructive={confirm}) expected {expected:?} got {decided:?}" + ); + // Compat invariant: allows() == (decide() != Deny) + assert_eq!( + policy.allows(risk, None), + !matches!(decided, PolicyAction::Deny), + "allows() must equal (decide() != Deny) for {mode:?}/{risk:?}/{confirm}" + ); + } + } + + #[test] + fn test_decide_connection_overrides() { + // allowlist excludes a connection → Deny + let allowlist = AiPolicy { + allowed_connection_ids: vec!["conn-1".into()], + ..AiPolicy::default() + }; + assert_eq!( + allowlist.decide(RiskLevel::Safe, Some("conn-2")), + PolicyAction::Deny + ); + assert_eq!( + allowlist.decide(RiskLevel::Safe, Some("conn-1")), + PolicyAction::Allow + ); + + // read-only override blocks Elevated for that connection + let override_policy = AiPolicy { + mode: McpPermissionMode::FullAccess, + connection_overrides: HashMap::from([( + "prod-es".into(), + ConnectionMcpOverride { + read_only: true, + allowed_actions: None, + }, + )]), + ..AiPolicy::default() + }; + assert_eq!( + override_policy.decide(RiskLevel::Elevated, Some("prod-es")), + PolicyAction::Deny + ); + assert_eq!( + override_policy.decide(RiskLevel::Safe, Some("prod-es")), + PolicyAction::Allow + ); + assert_eq!( + override_policy.decide(RiskLevel::Elevated, Some("staging")), + PolicyAction::Allow + ); + } + + #[test] + fn test_policy_action_serde_names() { + assert_eq!( + serde_json::to_value(PolicyAction::Allow).unwrap(), + json!("allow") + ); + assert_eq!( + serde_json::to_value(PolicyAction::Ask).unwrap(), + json!("ask") + ); + assert_eq!( + serde_json::to_value(PolicyAction::Deny).unwrap(), + json!("deny") + ); + } + + #[test] + fn test_connection_action_allowlist_blocks_risks() { + // Read-only action list: only Safe capabilities pass + let policy = AiPolicy { + mode: McpPermissionMode::FullAccess, + connection_overrides: HashMap::from([( + "prod".into(), + ConnectionMcpOverride { + read_only: false, + allowed_actions: Some(vec![McpAction::Read]), + }, + )]), + ..AiPolicy::default() + }; + assert_eq!( + policy.decide(RiskLevel::Safe, Some("prod")), + PolicyAction::Allow + ); + assert_eq!( + policy.decide(RiskLevel::Elevated, Some("prod")), + PolicyAction::Deny + ); + assert_eq!( + policy.decide(RiskLevel::Destructive, Some("prod")), + PolicyAction::Deny + ); + } + + #[test] + fn test_connection_action_allowlist_read_write() { + let policy = AiPolicy { + mode: McpPermissionMode::FullAccess, + connection_overrides: HashMap::from([( + "prod".into(), + ConnectionMcpOverride { + read_only: false, + allowed_actions: Some(vec![McpAction::Read, McpAction::Write]), + }, + )]), + ..AiPolicy::default() + }; + assert_eq!( + policy.decide(RiskLevel::Safe, Some("prod")), + PolicyAction::Allow + ); + assert_eq!( + policy.decide(RiskLevel::Elevated, Some("prod")), + PolicyAction::Allow + ); + assert_eq!( + policy.decide(RiskLevel::Destructive, Some("prod")), + PolicyAction::Deny + ); + } + + #[test] + fn test_read_only_legacy_still_works_without_actions() { + // Legacy read_only override (no allowed_actions) keeps old behavior + let policy = AiPolicy { + mode: McpPermissionMode::FullAccess, + connection_overrides: HashMap::from([( + "prod".into(), + ConnectionMcpOverride { + read_only: true, + allowed_actions: None, + }, + )]), + ..AiPolicy::default() + }; + assert_eq!( + policy.decide(RiskLevel::Safe, Some("prod")), + PolicyAction::Allow + ); + assert_eq!( + policy.decide(RiskLevel::Elevated, Some("prod")), + PolicyAction::Deny + ); + assert!(policy.is_connection_read_only("prod")); + } + + #[test] + fn test_mcp_action_serde_names() { + assert_eq!( + serde_json::to_value(McpAction::Read).unwrap(), + json!("read") + ); + assert_eq!( + serde_json::to_value(McpAction::Write).unwrap(), + json!("write") + ); + assert_eq!( + serde_json::to_value(McpAction::Delete).unwrap(), + json!("delete") + ); + assert_eq!(McpAction::from_risk(RiskLevel::Safe), McpAction::Read); + assert_eq!(McpAction::from_risk(RiskLevel::Elevated), McpAction::Write); + assert_eq!( + McpAction::from_risk(RiskLevel::Destructive), + McpAction::Delete + ); + } +} diff --git a/src-tauri/src/ai_permissions/scope.rs b/src-tauri/src/ai_permissions/scope.rs new file mode 100644 index 00000000..ae0a6da0 --- /dev/null +++ b/src-tauri/src/ai_permissions/scope.rs @@ -0,0 +1,288 @@ +use super::PermissionRules; +use data_studio_agent::{capabilities::types::RiskLevel, storage::db::AgentDb}; +use rusqlite::OptionalExtension; +use serde_json::Value; +use std::collections::HashMap; + +pub(crate) struct AgentAccess { + pub session_id: String, + pub connections: HashMap, + pub studio: bool, + pub db: AgentDb, +} + +pub(super) struct SessionAccess { + pub restrictions: HashMap, + pub(super) mode: String, + pub(super) rule: Option, +} + +impl AgentAccess { + pub(super) fn current(&self, tool: &str) -> Result { + let db = self.db.0.lock().map_err(|e| e.to_string())?; + let (sources, mode): (String, String) = db + .query_row( + "SELECT sources, permissions_mode FROM agent_sessions WHERE id = ?1", + [&self.session_id], + |row| Ok((row.get(0)?, row.get(1)?)), + ) + .map_err(|e| format!("Cannot load current conversation permissions: {e}"))?; + let sources: Vec = serde_json::from_str(&sources).map_err(|e| e.to_string())?; + let rule = db + .query_row( + "SELECT action FROM confirmation_rules WHERE session_id = ?1 AND tool_name = ?2", + [&self.session_id, tool], + |row| row.get(0), + ) + .optional() + .map_err(|e| e.to_string())?; + let mut restrictions = HashMap::new(); + for config in self.connections.values() { + let Some(id) = config.get("connectionId").and_then(Value::as_str) else { + continue; + }; + let permission = if self.studio { + let Some(source_id) = config.get("sourceId").and_then(Value::as_str) else { + continue; + }; + let Some(source) = sources.iter().find(|source| { + source.get("sourceId").and_then(Value::as_str) == Some(source_id) + && source.get("detached").and_then(Value::as_bool) != Some(true) + && source.get("kind").and_then(Value::as_str) == Some("database") + }) else { + continue; + }; + let attached_id: Option = db.query_row( + "SELECT CAST(connection_id AS TEXT) FROM attached_sources WHERE id = ?1 AND kind = 'database'", + [source_id], |row| row.get(0), + ).optional().map_err(|e| e.to_string())?.flatten(); + if attached_id.as_deref() != Some(id) { + continue; + } + if source.get("permissionsMode").and_then(Value::as_str) == Some("custom") { + serde_json::from_value( + source + .get("permissions") + .cloned() + .ok_or("Missing source permissions")?, + ) + .map_err(|e| e.to_string())? + } else { + PermissionRules::explicit(true, true, true, true, true) + } + } else { + PermissionRules::explicit(true, true, true, true, true) + }; + // A connection attached more than once must satisfy every attachment's limits. + restrictions + .entry(id.to_owned()) + .and_modify(|existing: &mut PermissionRules| { + *existing = intersect(existing, &permission); + }) + .or_insert(permission); + } + Ok(SessionAccess { + restrictions, + mode, + rule, + }) + } +} + +fn intersect(a: &PermissionRules, b: &PermissionRules) -> PermissionRules { + let combine = |a: Option, b: Option| { + if a == Some(false) || b == Some(false) { + Some(false) + } else { + a.or(b) + } + }; + PermissionRules { + read: combine(a.read, b.read), + create: combine(a.create, b.create), + update: combine(a.update, b.update), + delete: combine(a.delete, b.delete), + // Custom conversation permissions must explicitly grant manage. + manage: Some(a.manage == Some(true) && b.manage == Some(true)), + } +} + +impl SessionAccess { + pub fn confirmation_required(&self, risk: RiskLevel) -> Result { + match self.rule.as_deref() { + Some("deny") => Err("Tool call denied by conversation rule".into()), + Some("allow") => Ok(false), + _ => Ok(match risk { + RiskLevel::Safe => false, + RiskLevel::Destructive => true, + RiskLevel::Elevated => self.mode != "Auto", + }), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + use std::sync::{Arc, Mutex}; + + fn setup(studio: bool) -> AgentAccess { + let db = AgentDb(Arc::new(Mutex::new( + rusqlite::Connection::open_in_memory().unwrap(), + ))); + data_studio_agent::storage::db::migrate(&db).unwrap(); + db.0.lock().unwrap().execute_batch("INSERT INTO agent_sessions (id,title,created_at,updated_at,permissions_mode) VALUES ('chat','Chat',0,0,'Auto'); INSERT INTO attached_sources (id,kind,connection_id,created_at,updated_at) VALUES ('source','database',42,0,0);").unwrap(); + let access = AgentAccess { + session_id: "chat".into(), + studio, + db, + connections: HashMap::from([( + "alias".into(), + json!({ + "connectionId":"42", "sourceId":"source", "permissions":{"read":true} + }), + )]), + }; + sources( + &access, + json!([{"sourceId":"source","kind":"database","permissionsMode":"inherit"}]), + ); + access + } + + fn sources(access: &AgentAccess, value: Value) { + access + .db + .0 + .lock() + .unwrap() + .execute( + "UPDATE agent_sessions SET sources=?1 WHERE id='chat'", + [value.to_string()], + ) + .unwrap(); + } + + #[test] + fn reloads_source_restrictions_instead_of_the_run_snapshot() { + let access = setup(true); + assert_eq!( + access.current("tool").unwrap().restrictions["42"].read, + Some(true) + ); + sources( + &access, + json!([{"sourceId":"source","kind":"database","permissionsMode":"custom","permissions":{"read":false,"update":true}}]), + ); + let current = access.current("tool").unwrap(); + assert_eq!(current.restrictions["42"].read, Some(false)); + assert_eq!(current.restrictions["42"].update, Some(true)); + assert_eq!(current.restrictions["42"].manage, None); + } + + #[test] + fn detaching_or_deleting_a_source_revokes_queued_access() { + let access = setup(true); + sources( + &access, + json!([{"sourceId":"source","kind":"database","detached":true}]), + ); + assert!(access.current("tool").unwrap().restrictions.is_empty()); + sources( + &access, + json!([{"sourceId":"source","kind":"database","permissionsMode":"inherit"}]), + ); + access + .db + .0 + .lock() + .unwrap() + .execute("DELETE FROM attached_sources", []) + .unwrap(); + assert!(access.current("tool").unwrap().restrictions.is_empty()); + } + + #[test] + fn rebinding_a_source_does_not_grant_access_to_the_old_target() { + let access = setup(true); + access + .db + .0 + .lock() + .unwrap() + .execute("UPDATE attached_sources SET connection_id=99", []) + .unwrap(); + assert!(access.current("tool").unwrap().restrictions.is_empty()); + } + + #[test] + fn sidebar_scope_stays_bound_without_requiring_studio_attachments() { + let access = setup(false); + sources(&access, json!([])); + let scope = access.current("tool").unwrap(); + assert!(scope.restrictions.contains_key("42")); + assert!(!scope.restrictions.contains_key("99")); + } + + #[test] + fn confirmation_mode_and_rules_are_reloaded_for_queued_calls() { + let access = setup(false); + assert!(!access + .current("tool") + .unwrap() + .confirmation_required(RiskLevel::Elevated) + .unwrap()); + access + .db + .0 + .lock() + .unwrap() + .execute("UPDATE agent_sessions SET permissions_mode='Ask'", []) + .unwrap(); + assert!(access + .current("tool") + .unwrap() + .confirmation_required(RiskLevel::Elevated) + .unwrap()); + access.db.0.lock().unwrap().execute("INSERT INTO confirmation_rules (id,session_id,tool_name,action,created_at) VALUES ('rule','chat','tool','allow',0)", []).unwrap(); + assert!(!access + .current("tool") + .unwrap() + .confirmation_required(RiskLevel::Destructive) + .unwrap()); + access + .db + .0 + .lock() + .unwrap() + .execute("UPDATE confirmation_rules SET action='deny'", []) + .unwrap(); + assert!(access + .current("tool") + .unwrap() + .confirmation_required(RiskLevel::Safe) + .is_err()); + } + + #[test] + fn missing_or_malformed_session_state_fails_closed() { + let access = setup(true); + access + .db + .0 + .lock() + .unwrap() + .execute("UPDATE agent_sessions SET sources='invalid'", []) + .unwrap(); + assert!(access.current("tool").is_err()); + access + .db + .0 + .lock() + .unwrap() + .execute("DELETE FROM agent_sessions", []) + .unwrap(); + assert!(access.current("tool").is_err()); + } +} diff --git a/src-tauri/src/capabilities/commands.rs b/src-tauri/src/capabilities/commands.rs index af16e0a9..6fb706a2 100644 --- a/src-tauri/src/capabilities/commands.rs +++ b/src-tauri/src/capabilities/commands.rs @@ -1,27 +1,6 @@ use data_studio_agent::capabilities::registry; use data_studio_agent::capabilities::types::Capability; use serde_json::{json, Value}; -use tauri::AppHandle; -use tauri::Manager; - -#[tauri::command] -pub async fn invoke_capability( - name: String, - args: Value, - connection_id: Option, - app: AppHandle, -) -> Result { - // The same gate the agent loop this capability surface serves is under. - crate::entitlement::ensure_local_ultimate( - &app.state::(), - "AI", - )?; - let config = match connection_id { - Some(ref id) => Some(resolve_connection_config(&app, id).await?), - None => None, - }; - registry::invoke_capability_inner(&name, args, config).await -} #[tauri::command] pub async fn get_available_tools(source_kinds: Option>) -> Result { @@ -60,10 +39,3 @@ fn to_metadata(cap: &Capability) -> Value { "requiredPermission": cap.required_permission }) } - -async fn resolve_connection_config(app: &AppHandle, connection_id: &str) -> Result { - let state: tauri::State<'_, crate::state::AppState> = app.state(); - state.ensure_connection(connection_id).await?; - - Ok(json!({ "connectionId": connection_id })) -} 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..f14cc5c7 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -19,6 +19,7 @@ pub mod menu; // Agent / AI modules pub mod agent; pub mod agent_adapters; +pub(crate) mod ai_permissions; pub mod capabilities; pub mod common; pub mod db; @@ -41,7 +42,6 @@ pub static APP_HANDLE: OnceLock = OnceLock::new(); pub static GUARDIAN: OnceLock> = OnceLock::new(); -use agent::executor::SqlKitToolExecutor; use agent::query_history::{add_query_history_entry, load_query_history}; use agent::session_store::{ clear_agent_session_messages, clear_session_confirmation_rules, create_agent_session, @@ -55,7 +55,7 @@ use agent_adapters::{ get_all_tools, get_tool_full_result, list_llm_models, run_agent_loop, run_agent_step, validate_llm_config, }; -use capabilities::commands::{get_available_tools, invoke_capability}; +use capabilities::commands::get_available_tools; use data_studio_agent as lib; use data_studio_agent::storage; @@ -189,8 +189,6 @@ pub fn run() { let cancel_map: lib::traits::CancelMap = Arc::new(Mutex::new(HashMap::new())); app.manage(confirm_map); app.manage(cancel_map); - let executor: Arc = Arc::new(SqlKitToolExecutor); - app.manage(executor); app.manage(crate::mcp_bridge::McpServerHandle::new()); { @@ -268,6 +266,9 @@ pub fn run() { Ok(()) }) .invoke_handler(tauri::generate_handler![ + crate::ai_permissions::get_ai_permissions, + ai_permissions::invocation::get_ai_context, + crate::ai_permissions::save_ai_permissions, // SSH profile management + host-key store crate::ssh::commands::list_ssh_profiles, crate::ssh::commands::save_ssh_profile, @@ -277,7 +278,6 @@ pub fn run() { crate::ssh::commands::unpin_ssh_host, crate::common::http_client::detect_system_proxy, // Agent commands - invoke_capability, get_available_tools, run_agent_step, validate_llm_config, diff --git a/src-tauri/src/mcp_bridge.rs b/src-tauri/src/mcp_bridge.rs index 6caef5b7..725469e6 100644 --- a/src-tauri/src/mcp_bridge.rs +++ b/src-tauri/src/mcp_bridge.rs @@ -2,12 +2,14 @@ use std::path::{Path, PathBuf}; use std::sync::Arc; use std::sync::Mutex; +use crate::ai_permissions::{ + load_permission_policy, read_permission_config, AiPolicy, AI_POLICY_LOCK, +}; use axum::extract::State; use axum::http::StatusCode; use axum::response::{IntoResponse, Response}; use axum::routing::{get, post}; use axum::Json; -use data_studio_agent::capabilities::permissions::McpPolicy; use data_studio_agent::capabilities::registry; use data_studio_agent::capabilities::types::{Capability, SourceKind}; use serde::{Deserialize, Serialize}; @@ -24,9 +26,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 +33,6 @@ impl McpServerHandle { Self { shutdown_tx: Mutex::new(None), server_task: Mutex::new(None), - policy: Arc::new(tokio::sync::RwLock::new(McpPolicy::default())), } } } @@ -50,7 +48,7 @@ pub struct McpConfig { #[serde(default = "default_auto_start")] pub auto_start: bool, #[serde(default)] - pub policy: McpPolicy, + pub policy: AiPolicy, } fn default_auto_start() -> bool { @@ -78,7 +76,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}") + }) } } @@ -87,7 +91,7 @@ impl Default for McpConfig { Self { port: None, auto_start: true, - policy: McpPolicy::default(), + policy: AiPolicy::default(), } } } @@ -138,21 +142,20 @@ struct BridgeState { handle: AppHandle, app_name: &'static str, app_data_dir: PathBuf, - policy: Arc>, } // --------------------------------------------------------------------------- // Handlers // --------------------------------------------------------------------------- -fn tools_payload(policy: &McpPolicy) -> Value { +fn tools_payload(policy: &AiPolicy) -> Value { let reg = registry::registry(); let caps = reg.agent_tools(); let notice = policy.policy_notice(); 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 +204,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() @@ -255,14 +267,17 @@ 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 { +fn list_allowed_connections(policy: &AiPolicy) -> Value { + if policy.permission_defaults.is_some() { + return filter_ai_connections(list_connections(), policy); + } if policy.allowed_connection_ids.is_empty() { return list_connections(); } filter_connections_by_allowlist(list_connections(), policy) } -fn filter_connections_by_allowlist(connections: Value, policy: &McpPolicy) -> Value { +fn filter_connections_by_allowlist(connections: Value, policy: &AiPolicy) -> Value { let filtered: Vec = connections .as_array() .map(|arr| { @@ -284,7 +299,7 @@ fn filter_connections_by_allowlist(connections: Value, policy: &McpPolicy) -> Va /// the global mode + confirm_destructive gate, not per-connection overrides. fn check_policy( cap: &Capability, - policy: &McpPolicy, + policy: &AiPolicy, connection_id: Option<&str>, ) -> Result<(), String> { if policy.allows(cap.risk_level, connection_id) { @@ -303,8 +318,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: &AiPolicy, + 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(), @@ -320,15 +342,35 @@ async fn handle_invoke( if let Some(reason) = entitlement_denied(&state) { return Json(InvokeResponse::error(403, reason)); } - // 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 { +async fn invoke_with_policy(policy: &AiPolicy, 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 +379,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,24 +391,30 @@ 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); } - let config = match payload.connection_id { - Some(ref id) => match resolve_connection(id).await { - Ok(cfg) => Some(cfg), - Err(e) => return InvokeResponse::error(400, e), - }, - None => None, - }; - - match registry::invoke_capability_inner(&payload.name, payload.args, config).await { + let config = payload + .connection_id + .as_ref() + .map(|id| json!({"connectionId": id})) + .unwrap_or(Value::Null); + match crate::ai_permissions::invoke_ai(&payload.name, &payload.args, &config, None).await { Ok(data) => match serde_json::from_str::(&data) { Ok(parsed) => InvokeResponse::ok(parsed), Err(_) => InvokeResponse::ok(json!({"result": data})), }, - Err(msg) => InvokeResponse::error(400, msg), + Err(data_studio_agent::tool_executor::ToolExecutionError::Denied(msg)) => { + InvokeResponse::error(403, msg) + } + Err(error) => InvokeResponse::error(400, error.to_string()), } } @@ -507,14 +557,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 +564,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 +609,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() @@ -615,13 +655,6 @@ fn to_metadata(cap: &Capability) -> Value { }) } -async fn resolve_connection(connection_id: &str) -> Result { - // Shared resolver: connects saved-but-inactive connections in-process; - // credentials never leave the app. McpPolicy auth already ran upstream. - crate::capabilities::sql::resolve_adapter(connection_id).await?; - Ok(json!({ "connectionId": connection_id })) -} - // --------------------------------------------------------------------------- // Tauri commands // --------------------------------------------------------------------------- @@ -633,7 +666,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!({ @@ -652,7 +685,7 @@ pub async fn save_mcp_config( port: Option, auto_start: bool, restart: Option, - policy: Option, + policy: Option, app: AppHandle, ) -> Result { use tauri::Manager; @@ -664,14 +697,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); @@ -704,7 +740,7 @@ pub async fn save_mcp_config( /// Persist the MCP policy and hot-swap it into the running bridge without a /// restart, so permission changes never interrupt in-flight LLM requests. #[tauri::command] -pub async fn save_mcp_policy(policy: McpPolicy, app: AppHandle) -> Result { +pub async fn save_mcp_policy(policy: AiPolicy, app: AppHandle) -> Result { use tauri::Manager; crate::entitlement::ensure_local_ultimate( &app.state::(), @@ -716,20 +752,37 @@ 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()) } +fn filter_ai_connections(connections: Value, policy: &AiPolicy) -> 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::>()) +} + #[cfg(test)] mod tests { use super::*; - use data_studio_agent::capabilities::permissions::McpPermissionMode; + use crate::ai_permissions::McpPermissionMode; use data_studio_agent::capabilities::types::RiskLevel; fn init_registry_for_tests() { @@ -760,7 +813,7 @@ mod tests { let cfg = McpConfig::default(); assert_eq!(cfg.port, None); assert!(cfg.auto_start); - assert_eq!(cfg.policy, McpPolicy::default()); + assert_eq!(cfg.policy, AiPolicy::default()); } #[test] @@ -775,10 +828,10 @@ mod tests { let cfg = McpConfig { port: Some(9444), auto_start: false, - policy: McpPolicy { + policy: AiPolicy { mode: McpPermissionMode::DataReadWrite, confirm_destructive: false, - ..McpPolicy::default() + ..AiPolicy::default() }, }; cfg.save(&dir).unwrap(); @@ -803,7 +856,7 @@ mod tests { let cfg = McpConfig::load(&dir); assert_eq!(cfg.port, None); assert!(cfg.auto_start); - assert_eq!(cfg.policy, McpPolicy::default()); + assert_eq!(cfg.policy, AiPolicy::default()); std::fs::remove_dir_all(&dir).unwrap(); } @@ -873,7 +926,7 @@ mod tests { #[test] fn test_tools_payload_is_flat_with_metadata() { init_registry_for_tests(); - let v = tools_payload(&McpPolicy::default()); + let v = tools_payload(&AiPolicy::default()); let tools = v["tools"].as_array().unwrap(); assert!(!tools.is_empty()); for t in tools { @@ -899,7 +952,7 @@ mod tests { return; } - let v = tools_payload(&McpPolicy::default()); + let v = tools_payload(&AiPolicy::default()); let names: std::collections::HashSet<&str> = v["tools"] .as_array() .unwrap() @@ -916,9 +969,9 @@ mod tests { ); } - let full = McpPolicy { + let full = AiPolicy { mode: McpPermissionMode::FullAccess, - ..McpPolicy::default() + ..AiPolicy::default() }; let v_full = tools_payload(&full); assert_eq!(v_full["tools"].as_array().unwrap().len(), caps.len()); @@ -928,7 +981,7 @@ mod tests { fn test_tools_payload_appends_policy_notice_when_gated() { init_registry_for_tests(); // Default DataReadWrite gates Destructive → notice appended - let v = tools_payload(&McpPolicy::default()); + let v = tools_payload(&AiPolicy::default()); let tools = v["tools"].as_array().unwrap(); assert!(!tools.is_empty()); for t in tools { @@ -940,9 +993,9 @@ mod tests { assert!(desc.contains("destructive operations")); } // FullAccess + confirm on → no gate → no notice - let full = McpPolicy { + let full = AiPolicy { mode: McpPermissionMode::FullAccess, - ..McpPolicy::default() + ..AiPolicy::default() }; let v_full = tools_payload(&full); for t in v_full["tools"].as_array().unwrap() { @@ -967,15 +1020,15 @@ mod tests { return; }; - let err = check_policy(risky, &McpPolicy::default(), None).unwrap_err(); + let err = check_policy(risky, &AiPolicy::default(), None).unwrap_err(); assert!(err.contains("blocked by MCP policy")); // Error carries actionable guidance for the agent to relay assert!(err.contains("requires FullAccess permission mode")); - let full_no_confirm = McpPolicy { + let full_no_confirm = AiPolicy { mode: McpPermissionMode::FullAccess, confirm_destructive: false, - ..McpPolicy::default() + ..AiPolicy::default() }; if let Some(destructive) = caps .iter() @@ -995,7 +1048,7 @@ mod tests { }; let rt = tokio::runtime::Runtime::new().unwrap(); - let resp = rt.block_on(invoke_with_policy(&McpPolicy::default(), req)); + let resp = rt.block_on(invoke_with_policy(&AiPolicy::default(), req)); assert_eq!(resp.status, 404); assert!(resp.message.unwrap().contains("Unknown capability")); @@ -1018,16 +1071,17 @@ mod tests { return; }; - let allowlist = McpPolicy { + let allowlist = AiPolicy { allowed_connection_ids: vec!["conn-1".into()], - ..McpPolicy::default() + ..AiPolicy::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, &AiPolicy::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] @@ -1046,10 +1100,10 @@ mod tests { args: json!({}), connection_id: None, }; - let allowlist = McpPolicy { + let allowlist = AiPolicy { mode: McpPermissionMode::FullAccess, allowed_connection_ids: vec!["conn-1".into()], - ..McpPolicy::default() + ..AiPolicy::default() }; let rt = tokio::runtime::Runtime::new().unwrap(); @@ -1068,13 +1122,13 @@ mod tests { ]); // Empty allowlist matches nothing (the caller short-circuits on it) - let none = filter_connections_by_allowlist(conns.clone(), &McpPolicy::default()); + let none = filter_connections_by_allowlist(conns.clone(), &AiPolicy::default()); assert_eq!(none.as_array().unwrap().len(), 0); // Non-empty allowlist → only listed connections survive - let allowlist = McpPolicy { + let allowlist = AiPolicy { allowed_connection_ids: vec!["conn-2".into()], - ..McpPolicy::default() + ..AiPolicy::default() }; let filtered = filter_connections_by_allowlist(conns, &allowlist); let arr = filtered.as_array().unwrap(); @@ -1102,14 +1156,14 @@ mod tests { }; let rt = tokio::runtime::Runtime::new().unwrap(); - let resp = rt.block_on(invoke_with_policy(&McpPolicy::default(), req)); + let resp = rt.block_on(invoke_with_policy(&AiPolicy::default(), req)); assert_eq!(resp.status, 403); assert!(resp.message.unwrap().contains("blocked by MCP policy")); } #[test] - fn test_handle_invoke_full_access_allows_risky_but_fails_invoke() { + fn test_handle_invoke_without_host_fails_closed() { init_registry_for_tests(); let tools = registry::registry().agent_tools(); let Some(risky) = tools @@ -1125,16 +1179,15 @@ mod tests { connection_id: None, }; - let full = McpPolicy { + let full = AiPolicy { mode: McpPermissionMode::FullAccess, - ..McpPolicy::default() + ..AiPolicy::default() }; let rt = tokio::runtime::Runtime::new().unwrap(); let resp = rt.block_on(invoke_with_policy(&full, req)); - // Policy passes; with no connection config the capability itself fails, - // proving execution reached the invoke path. - assert_eq!(resp.status, 400); + assert_eq!(resp.status, 403); + assert!(resp.message.unwrap().contains("App not initialized")); } #[test] @@ -1165,4 +1218,67 @@ mod tests { ); } } + #[test] + fn ai_permission_roundtrip_preserves_partial_decisions_and_service_settings() { + use crate::ai_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 = AiPolicy { + 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 crate::ai_permissions::PermissionRules; + let mut policy = AiPolicy { + 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!([])); + } } diff --git a/src/__tests__/agentRuntimePermissions.test.ts b/src/__tests__/agentRuntimePermissions.test.ts new file mode 100644 index 00000000..b5dc92fd --- /dev/null +++ b/src/__tests__/agentRuntimePermissions.test.ts @@ -0,0 +1,115 @@ +import type { AgentToolCall } from '@/store/dataStudioStore' +import { disposeAgentRuntime, initAgentRuntime } from '@/composables/agentRuntime' + +const mockHandlers: Record void> = {} +const mockConfirm = jest.fn().mockResolvedValue(undefined) +const mockSession = { + id: 'chat', + permissionsMode: 'Auto', + messages: [{ id: 'assistant', role: 'assistant', status: 'streaming', toolCalls: [] as AgentToolCall[] }], +} +const mockStore = { + sessions: [mockSession], + removePreparingPlaceholder: jest.fn(), + setSessionActiveTool: jest.fn(), + setSessionStatus: jest.fn(), + updateToolCallStatus: jest.fn(), + findConfirmationRule: jest.fn(), + setMessageToolCalls: jest.fn((messageId: string, calls: AgentToolCall[], _sessionId: string) => { + mockSession.messages = mockSession.messages.map(message => message.id === messageId ? { ...message, toolCalls: calls } : message) + }), +} + +jest.mock('@/store/dataStudioStore', () => ({ useDataStudioStore: () => mockStore })) +jest.mock('@/datasources/agentApi', () => { + const register = (name: string) => async (handler: (payload: any) => void) => { + mockHandlers[name] = handler + return () => undefined + } + const api = { + onAgentContextUsage: register('onAgentContextUsage'), + onAgentLoopCompacting: register('onAgentLoopCompacting'), + onAgentLoopDelta: register('onAgentLoopDelta'), + onAgentLoopDone: register('onAgentLoopDone'), + onAgentLoopError: register('onAgentLoopError'), + onAgentLoopIteration: register('onAgentLoopIteration'), + onAgentLoopStepDone: register('onAgentLoopStepDone'), + onAgentLoopStopped: register('onAgentLoopStopped'), + onAgentLoopSummaryInjected: register('onAgentLoopSummaryInjected'), + onAgentLoopThinkingDelta: register('onAgentLoopThinkingDelta'), + onAgentLoopToolCall: register('onAgentLoopToolCall'), + onAgentLoopToolResult: register('onAgentLoopToolResult'), + onAgentLoopToolRetry: register('onAgentLoopToolRetry'), + onAgentLoopWaitingLlm: register('onAgentLoopWaitingLlm'), + onAgentLoopWarning: register('onAgentLoopWarning'), + confirmToolCall: (...args: unknown[]) => mockConfirm(...args), + } + return { ...api, agentApi: api } +}) + +const call = (confirmation_required?: boolean) => ({ + session_id: 'chat', + tool_call_id: 'call', + tool_name: 'update', + arguments: { value: 1 }, + risk_level: 'elevated', + confirmation_required, +}) + +describe('agent permission events', () => { + beforeEach(async () => { + jest.useFakeTimers() + mockSession.messages = [{ id: 'assistant', role: 'assistant', status: 'streaming', toolCalls: [] }] + jest.clearAllMocks() + mockStore.findConfirmationRule.mockReturnValue(undefined) + await initAgentRuntime() + }) + afterEach(() => { + disposeAgentRuntime() + jest.useRealTimers() + }) + + it('honors current host confirmation requirements over stale frontend Auto and allow rules', () => { + mockStore.findConfirmationRule.mockReturnValue({ action: 'allow' }) + mockHandlers.onAgentLoopToolCall(call(true)) + expect(mockConfirm).not.toHaveBeenCalled() + expect(mockSession.messages[0].toolCalls[0].status).toBe('pending') + expect(mockStore.setSessionStatus).toHaveBeenCalledWith('chat', 'waiting_confirmation') + }) + + it('reopens one queued call for approval without duplicating it or keeping its execution timer', () => { + mockHandlers.onAgentLoopToolCall(call(false)) + expect(mockConfirm).toHaveBeenCalledWith('call', true) + mockConfirm.mockClear() + mockHandlers.onAgentLoopToolCall(call(true)) + expect(mockSession.messages[0].toolCalls).toHaveLength(1) + expect(mockSession.messages[0].toolCalls[0].requiresConfirmation).toBe(true) + jest.advanceTimersByTime(61_000) + expect(mockConfirm).not.toHaveBeenCalled() + expect(mockStore.updateToolCallStatus).not.toHaveBeenCalled() + }) + + it('renders a preflight denial without requesting confirmation', () => { + mockHandlers.onAgentLoopToolResult({ + ...call(), + denied: true, + error: true, + envelope: { summary: 'Read access denied' }, + }) + expect(mockConfirm).not.toHaveBeenCalled() + expect(mockSession.messages[0].toolCalls).toHaveLength(1) + expect(mockSession.messages[0].toolCalls[0]).toMatchObject({ status: 'denied', result: 'Read access denied', requiresConfirmation: false }) + }) + + it('reports execution-time revocation as denial', () => { + mockHandlers.onAgentLoopToolCall(call(false)) + mockHandlers.onAgentLoopToolResult({ ...call(), denied: true, error: true, envelope: { summary: 'Permission revoked' } }) + expect(mockStore.updateToolCallStatus).toHaveBeenCalledWith('assistant', 'call', 'denied', 'Permission revoked', undefined, 'chat') + }) + + it('retains existing confirmation behavior for runtimes without host metadata', () => { + mockHandlers.onAgentLoopToolCall({ ...call(), risk_level: 'destructive' }) + expect(mockConfirm).not.toHaveBeenCalled() + expect(mockSession.messages[0].toolCalls[0].requiresConfirmation).toBe(true) + }) +}) 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, confirmation_required }) => { + clearToolTimeout(tool_call_id) const store = useDataStudioStore() store.removePreparingPlaceholder() store.setSessionActiveTool(session_id, tool_name) @@ -157,8 +158,8 @@ async function initAgentRuntime(): Promise { return const runtime = getSessionRuntime(session_id) - const riskLevel = (runtime.metadata?.[tool_name]?.riskLevel as AgentToolCall['riskLevel']) ?? 'elevated' - const needsConfirmation = shouldRequireConfirmation(session, tool_name, riskLevel) + const riskLevel = risk_level ?? (runtime.metadata?.[tool_name]?.riskLevel as AgentToolCall['riskLevel']) ?? 'elevated' + const needsConfirmation = confirmation_required ?? shouldRequireConfirmation(session, tool_name, riskLevel) const denied = isDeniedByRule(session_id, tool_name) const toolCall: AgentToolCall = { @@ -170,16 +171,17 @@ async function initAgentRuntime(): Promise { requiresConfirmation: needsConfirmation, } + const existingAssistant = session.messages.find(message => message.toolCalls?.some(call => call.id === tool_call_id)) const streamingAssistant = [...session.messages] .reverse() .find(message => message.role === 'assistant' && message.status === 'streaming') const lastAssistant - = streamingAssistant + = existingAssistant ?? streamingAssistant ?? [...session.messages].reverse().find(message => message.role === 'assistant') if (lastAssistant) { store.setMessageToolCalls(lastAssistant.id, [ - ...(lastAssistant.toolCalls ?? []), + ...(lastAssistant.toolCalls ?? []).filter(call => call.id !== tool_call_id), toolCall, ], session_id) } @@ -213,7 +215,7 @@ async function initAgentRuntime(): Promise { ) unlisteners.push( - await agentApi.onAgentLoopToolResult(({ session_id, tool_call_id, envelope, error }) => { + await agentApi.onAgentLoopToolResult(({ session_id, tool_call_id, envelope, error, denied, tool_name, arguments: args }) => { clearToolTimeout(tool_call_id) const store = useDataStudioStore() store.setSessionActiveTool(session_id, null) @@ -229,11 +231,27 @@ async function initAgentRuntime(): Promise { && message.toolCalls?.some(toolCall => toolCall.id === tool_call_id), ) + if (!assistantMsg && denied && tool_name) { + const target = [...session.messages].reverse().find(message => message.role === 'assistant') + if (target) { + const rejected: AgentToolCall = { + id: tool_call_id, + toolName: tool_name, + args: JSON.stringify(args ?? {}), + status: 'denied', + riskLevel: 'safe', + requiresConfirmation: false, + result: envelope.summary, + } + store.setMessageToolCalls(target.id, [...(target.toolCalls ?? []), rejected], session_id) + } + } + if (assistantMsg) { store.updateToolCallStatus( assistantMsg.id, tool_call_id, - error ? 'error' : 'done', + denied ? 'denied' : error ? 'error' : 'done', envelope.summary, envelope.metadata?.duration_ms, session_id, diff --git a/src/composables/useChatAgent.ts b/src/composables/useChatAgent.ts index 847de98a..400de5c5 100644 --- a/src/composables/useChatAgent.ts +++ b/src/composables/useChatAgent.ts @@ -368,13 +368,14 @@ export function useChatAgent(config: UseChatAgentConfig) { const connectionConfig = contextSnapshot?.connections ? Object.fromEntries( Object.entries(contextSnapshot.connections).flatMap(([alias, conn]) => { - const entry = { connectionId: String(conn.connectionId), dbType: conn.dbType, permissions: conn.permissions } + const entry = { sourceId: conn.sourceId, connectionId: String(conn.connectionId), dbType: conn.dbType } return [[String(conn.connectionId), entry], [alias, entry]] }), ) : {} const settings: Record = { + contextKind: config.feature, systemPrompt: assembledSystemPrompt, provider: modelConfig.provider.name, model: modelConfig.model.label, @@ -491,27 +492,15 @@ export function useChatAgent(config: UseChatAgentConfig) { return } - let schemaContext = '' - if (connections && Object.keys(connections).length > 0) { - const schemaParts: string[] = [] - for (const [alias, conn] of Object.entries(connections)) { - try { - const schemaResult = await agentApi.invokeCapability('sqlkit__get_schema', {}, String(conn.connectionId)) - if (schemaResult && schemaResult !== '[]' && !schemaResult.startsWith('[')) { - schemaParts.push(`-- Source: ${alias} (${conn.dbType})`) - schemaParts.push(schemaResult) - } - } - catch { - // Schema fetching is best-effort; continue without it - } - } - if (schemaParts.length > 0) { - schemaContext = schemaParts.join('\n') - } + const aiContext = await agentApi.getAiContext(sessionId, connections, config.feature) + const canReadContext = Object.values(connections).length > 0 + && Object.values(connections).every(conn => aiContext.readableConnectionIds.includes(String(conn.connectionId))) + const safeContext = { + ...context, + schema: canReadContext ? context?.schema : undefined, + activePanel: canReadContext ? context?.activePanel : undefined, } - - await runAgentLoop(sessionId, options.content, schemaContext, context ?? { connections: {} }) + await runAgentLoop(sessionId, options.content, aiContext.schema, safeContext) } catch (err) { localError.value = err instanceof Error ? err.message : String(err) diff --git a/src/composables/useDataStudioChatAgent.ts b/src/composables/useDataStudioChatAgent.ts index 7514d997..20f78985 100644 --- a/src/composables/useDataStudioChatAgent.ts +++ b/src/composables/useDataStudioChatAgent.ts @@ -93,9 +93,10 @@ export function useDataStudioChatAgent() { return acc acc[sessionSource.alias] = { + sourceId: sessionSource.sourceId, 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..e60c7f81 100644 --- a/src/datasources/agentApi.ts +++ b/src/datasources/agentApi.ts @@ -247,6 +247,8 @@ export const agentApi = { session_id: string tool_call_id: string tool_name: string + risk_level?: 'safe' | 'elevated' | 'destructive' + confirmation_required?: boolean arguments: unknown }) => void, ): Promise<() => void> => @@ -258,6 +260,9 @@ export const agentApi = { tool_call_id: string envelope: ToolEnvelope error?: boolean + denied?: boolean + tool_name?: string + arguments?: unknown }) => void, ): Promise<() => void> => listen('agent-loop-tool-result', e => handler(e.payload as any)), @@ -351,13 +356,8 @@ export const agentApi = { return parseStringResult(result) }, - // Capability invocation — returns raw string - invokeCapability: ( - name: string, - args: Record, - connectionId?: string, - ): Promise => - invoke('invoke_capability', { name, args, connectionId }), + getAiContext: (sessionId: string, connections: Record, contextKind: string): Promise<{ readableConnectionIds: string[], schema: string }> => + invoke('get_ai_context', { sessionId, connections, contextKind }), } // Standalone function exports diff --git a/src/lang/enUS.ts b/src/lang/enUS.ts index 4497cc1c..02403034 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: @@ -570,32 +611,7 @@ export const enUS = { starting: 'Starting...', autoStart: 'Auto-start', autoStartDesc: 'Start the MCP bridge automatically when app launches', - permissionMode: 'Permission Mode', - permissionModeDesc: 'Control what MCP clients can do through the bridge', - modeReadOnly: 'Read-only', - modeReadOnlyDesc: 'Queries only, no modifications', - modeDataReadWrite: 'Data read/write', - modeDataReadWriteDesc: 'CRUD allowed, no schema changes', - modeFullAccess: 'Full access', - modeFullAccessDesc: 'All operations permitted', - confirmDestructive: 'Confirm Destructive', - confirmDestructiveDesc: 'Require explicit confirmation before executing destructive operations (DROP, TRUNCATE, DELETE, etc.)', - allowlist: 'Connection Allowlist', - allowlistDesc: 'Select which connections MCP clients can access. Leave empty to allow all.', - allowlistEmpty: 'No connections selected — all connections are accessible', - connectionOverrides: 'Connection Overrides', - overrideReadOnly: 'Read-only', - overrideReadOnlyDesc: 'Force read-only access for this connection regardless of the global mode', saveFailed: 'Failed to save MCP settings', - connectionAccess: 'Connection Access', - connectionAccessDesc: - 'When enabled, only listed connections are allowed. Set per-connection actions (read/write/delete).', - connectionName: 'Connection', - connectionType: 'Type', - allowedActions: 'Allowed Actions', - actionRead: 'Read', - actionWrite: 'Write', - actionDelete: 'Delete', }, about: { title: 'About', @@ -1110,8 +1126,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..1e82690d 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: @@ -570,32 +611,7 @@ export const zhCN = { starting: '正在启动...', autoStart: '自动启动', autoStartDesc: '应用启动时自动启动 MCP 桥接', - permissionMode: '权限模式', - permissionModeDesc: '控制 MCP 客户端通过桥接可执行的操作', - modeReadOnly: '只读', - modeReadOnlyDesc: '仅允许查询，不允许修改', - modeDataReadWrite: '数据读写', - modeDataReadWriteDesc: '允许 CRUD，不允许 schema 变更', - modeFullAccess: '完全访问', - modeFullAccessDesc: '允许所有操作', - confirmDestructive: '确认破坏性操作', - confirmDestructiveDesc: '执行破坏性操作（DROP、TRUNCATE、DELETE 等）前需要明确确认', - allowlist: '连接白名单', - allowlistDesc: '选择 MCP 客户端可访问的连接。留空表示允许所有连接。', - allowlistEmpty: '未选择任何连接 — 所有连接都可访问', - connectionOverrides: '连接覆盖', - overrideReadOnly: '只读', - overrideReadOnlyDesc: '无论全局模式如何，强制该连接为只读访问', saveFailed: '保存 MCP 设置失败', - connectionAccess: '连接访问控制', - connectionAccessDesc: - '启用后仅允许列表中的连接。可为每个连接设置操作权限（读/写/删除）。', - connectionName: '连接', - connectionType: '类型', - allowedActions: '允许的操作', - actionRead: '读', - actionWrite: '写', - actionDelete: '删', }, about: { title: '关于', @@ -1110,8 +1126,8 @@ export const zhCN = { modifySource: { title: '修改数据源', accessPermissions: '访问权限', - inheritTitle: '从会话继承', - inheritDesc: '权限遵循会话级设置并自动更新。', + inheritTitle: '使用连接权限', + inheritDesc: '遵循连接当前的权限链，未单独设置的操作使用系统默认权限。', customTitle: '自定义', customDesc: '手动定义该数据源允许的操作。', read: '读取', 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') }}