diff --git a/crates/dbx-mcp/src/backend.rs b/crates/dbx-mcp/src/backend.rs index 9db1877986..3c937ce147 100644 --- a/crates/dbx-mcp/src/backend.rs +++ b/crates/dbx-mcp/src/backend.rs @@ -27,6 +27,7 @@ pub struct ConnectionSummary { pub host: String, pub port: u16, pub database: String, + pub group_path: Vec, } impl From<&ConnectionConfig> for ConnectionSummary { @@ -42,6 +43,74 @@ impl From<&ConnectionConfig> for ConnectionSummary { host: config.host.clone(), port: config.port, database: config.database.clone().unwrap_or_default(), + group_path: Vec::new(), + } + } +} + +#[derive(Deserialize)] +struct SidebarLayout { + #[serde(default)] + groups: Vec, + #[serde(default)] + order: Vec, +} + +#[derive(Deserialize)] +struct SidebarGroup { + id: String, + name: String, +} + +#[derive(Deserialize)] +#[serde(tag = "type")] +enum SidebarOrderEntry { + #[serde(rename = "group")] + Group { + id: String, + children: Option>, + #[serde(rename = "connectionIds")] + connection_ids: Option>, + }, + #[serde(rename = "connection")] + Connection { id: String }, +} + +fn connection_group_paths(layout: Value) -> HashMap> { + let Ok(layout) = serde_json::from_value::(layout) else { + return HashMap::new(); + }; + let groups = layout.groups.into_iter().map(|group| (group.id, group.name)).collect::>(); + let mut paths = HashMap::new(); + collect_connection_group_paths(&layout.order, &groups, &mut Vec::new(), &mut paths); + paths +} + +fn collect_connection_group_paths( + entries: &[SidebarOrderEntry], + groups: &HashMap, + path: &mut Vec, + paths: &mut HashMap>, +) { + for entry in entries { + match entry { + SidebarOrderEntry::Connection { id } => { + paths.insert(id.clone(), path.clone()); + } + SidebarOrderEntry::Group { id, children, connection_ids } => { + let Some(name) = groups.get(id) else { + continue; + }; + path.push(name.clone()); + if let Some(children) = children { + collect_connection_group_paths(children, groups, path, paths); + } else if let Some(connection_ids) = connection_ids { + for connection_id in connection_ids { + paths.insert(connection_id.clone(), path.clone()); + } + } + path.pop(); + } } } } @@ -92,6 +161,9 @@ pub trait DbxBackend: Send + Sync { async fn load_mcp_global_policy(&self) -> Result; async fn load_connections(&self) -> Result, String>; + async fn load_connection_group_paths(&self) -> Result>, String> { + Ok(HashMap::new()) + } async fn execute_agent_tool( &self, connection: &ConnectionConfig, @@ -407,6 +479,10 @@ impl DbxBackend for LocalBackend { Ok(configs) } + async fn load_connection_group_paths(&self) -> Result>, String> { + Ok(self.state.storage.load_sidebar_layout().await?.map(connection_group_paths).unwrap_or_default()) + } + async fn execute_agent_tool( &self, connection: &ConnectionConfig, @@ -579,6 +655,16 @@ impl DbxBackend for WebBackend { .map_err(|error| format!("Invalid connection list response: {error}")) } + async fn load_connection_group_paths(&self) -> Result>, String> { + let layout = self + .request(reqwest::Method::GET, "/api/layout/sidebar", None) + .await? + .json::>() + .await + .map_err(|error| format!("Invalid sidebar layout response: {error}"))?; + Ok(layout.map(connection_group_paths).unwrap_or_default()) + } + async fn execute_agent_tool( &self, connection: &ConnectionConfig, @@ -1534,6 +1620,86 @@ mod tests { assert!(parse_database_type("unknown").is_err()); } + #[test] + fn parses_nested_current_and_legacy_connection_group_paths() { + let paths = connection_group_paths(json!({ + "groups": [ + { "id": "project", "name": "Project" }, + { "id": "staging", "name": "Staging" }, + { "id": "legacy", "name": "Legacy" } + ], + "order": [ + { + "type": "group", + "id": "project", + "children": [ + { + "type": "group", + "id": "staging", + "children": [{ "type": "connection", "id": "nested" }] + }, + { "type": "connection", "id": "grouped" } + ] + }, + { "type": "group", "id": "legacy", "connectionIds": ["legacy-connection"] }, + { + "type": "group", + "id": "missing-group", + "children": [{ "type": "connection", "id": "dangling" }] + }, + { "type": "connection", "id": "root" } + ] + })); + + assert_eq!(paths.get("nested"), Some(&vec!["Project".to_string(), "Staging".to_string()])); + assert_eq!(paths.get("grouped"), Some(&vec!["Project".to_string()])); + assert_eq!(paths.get("legacy-connection"), Some(&vec!["Legacy".to_string()])); + assert_eq!(paths.get("root"), Some(&Vec::new())); + assert!(!paths.contains_key("dangling")); + } + + #[test] + fn malformed_sidebar_layout_has_no_group_paths() { + assert!(connection_group_paths(json!({ "groups": "invalid", "order": [] })).is_empty()); + } + + #[tokio::test] + async fn web_backend_loads_connection_group_paths_from_sidebar_layout() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = listener.local_addr().unwrap(); + let (request_sender, request_receiver) = mpsc::channel(); + let server = std::thread::spawn(move || { + let (mut stream, _) = listener.accept().unwrap(); + stream.set_read_timeout(Some(Duration::from_secs(5))).unwrap(); + let mut request = Vec::new(); + let mut buffer = [0_u8; 1024]; + loop { + let count = stream.read(&mut buffer).unwrap(); + request.extend_from_slice(&buffer[..count]); + if count == 0 || request.windows(4).any(|window| window == b"\r\n\r\n") { + break; + } + } + request_sender.send(String::from_utf8(request).unwrap().lines().next().unwrap().to_string()).unwrap(); + let body = r#"{"groups":[{"id":"project","name":"Project"}],"order":[{"type":"group","id":"project","children":[{"type":"connection","id":"web-db"}]}]}"#; + write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", + body.len(), + body + ) + .unwrap(); + }); + + let backend = WebBackend::new(format!("http://{address}"), String::new()).unwrap(); + backend.auth.lock().await.checked = true; + let paths = backend.load_connection_group_paths().await.unwrap(); + + server.join().unwrap(); + assert_eq!(request_receiver.recv().unwrap(), "GET /api/layout/sidebar HTTP/1.1"); + assert_eq!(paths.get("web-db"), Some(&vec!["Project".to_string()])); + } + #[test] fn format_query_result_appends_server_messages() { let mut result = query_result(Vec::new(), Vec::new(), 3); diff --git a/crates/dbx-mcp/src/server.rs b/crates/dbx-mcp/src/server.rs index 4b7a3d67c2..7a83b18f3d 100644 --- a/crates/dbx-mcp/src/server.rs +++ b/crates/dbx-mcp/src/server.rs @@ -242,7 +242,7 @@ impl DbxMcpServer { impl DbxMcpServer { #[tool( name = "dbx_list_connections", - description = "List database connections configured in DBX. Returns connection IDs, names, database types, endpoints, and selected databases." + description = "List database connections configured in DBX. Returns connection IDs, names, group paths, database types, endpoints, and selected databases." )] async fn list_connections( &self, @@ -251,7 +251,15 @@ impl DbxMcpServer { match self.load_scoped_connections().await { Ok(connections) if connections.is_empty() => text("No connections configured in DBX."), Ok(connections) => { - let rows = connections.iter().map(ConnectionSummary::from).collect::>(); + let group_paths = self.backend.load_connection_group_paths().await.unwrap_or_default(); + let rows = connections + .iter() + .map(|connection| { + let mut summary = ConnectionSummary::from(connection); + summary.group_path = group_paths.get(&connection.id).cloned().unwrap_or_default(); + summary + }) + .collect::>(); text(format_connections(&rows)) } Err(error) => backend_tool_error("CONNECTION_LOAD_ERROR", error), @@ -1126,13 +1134,15 @@ fn ambiguous_connections(name: &str, connections: &[dbx_core::models::connection } fn format_connections(connections: &[ConnectionSummary]) -> String { - let mut output = - String::from("| ID | Name | Type | Host | Port | Database |\n| --- | --- | --- | --- | --- | --- |"); + let mut output = String::from( + "| ID | Name | Group Path | Type | Host | Port | Database |\n| --- | --- | --- | --- | --- | --- | --- |", + ); for connection in connections { output.push_str(&format!( - "\n| {} | {} | {} | {} | {} | {} |", + "\n| {} | {} | {} | {} | {} | {} | {} |", escape_cell(&connection.id), escape_cell(&connection.name), + escape_cell(&connection.group_path.join(" / ")), escape_cell(&connection.db_type), escape_cell(&connection.host), connection.port, @@ -1346,9 +1356,11 @@ mod tests { host: "127.0.0.1".to_string(), port: 5432, database: "app".to_string(), + group_path: vec!["Project|A".to_string(), "Staging\nWest".to_string()], }]); assert!(output.contains("id\\|1")); assert!(output.contains("local pg")); + assert!(output.contains("Project\\|A / Staging West")); } #[test] diff --git a/crates/dbx-mcp/tests/local.rs b/crates/dbx-mcp/tests/local.rs index 1e2f76c324..9d3ab3d6c6 100644 --- a/crates/dbx-mcp/tests/local.rs +++ b/crates/dbx-mcp/tests/local.rs @@ -1,7 +1,7 @@ use std::{ffi::OsString, sync::Arc}; use dbx_core::{models::connection::ConnectionConfig, storage::Storage}; -use dbx_mcp::{DbxBackend, DbxMcpServer, LocalBackend, McpScope}; +use dbx_mcp::{DbxMcpServer, LocalBackend, McpScope}; use rmcp::{model::CallToolRequestParams, ServiceExt}; use serde_json::{json, Map, Value}; use tempfile::tempdir; @@ -47,6 +47,24 @@ async fn local_backend_reads_dbx_storage_without_desktop_process() { })) .expect("minimal connection config"); storage.save_connections(&[connection]).await.expect("save connection"); + storage + .save_sidebar_layout(&json!({ + "groups": [ + { "id": "project", "name": "Project", "collapsed": false }, + { "id": "environment", "name": "Staging", "collapsed": false } + ], + "order": [{ + "type": "group", + "id": "project", + "children": [{ + "type": "group", + "id": "environment", + "children": [{ "type": "connection", "id": "local-sqlite" }] + }] + }] + })) + .await + .expect("save sidebar layout"); let backend = Arc::new(LocalBackend::open(&db_path).await.expect("open local backend")); let server = DbxMcpServer::with_runtime_options(backend, McpScope::default(), false); @@ -61,6 +79,7 @@ async fn local_backend_reads_dbx_storage_without_desktop_process() { let text = result.content[0].as_text().expect("text response"); assert!(text.text.contains("offline-sqlite")); assert!(text.text.contains("local-sqlite")); + assert!(text.text.contains("Project / Staging")); client.cancel().await.expect("close client"); server_task.abort(); } diff --git a/crates/dbx-mcp/tests/protocol.rs b/crates/dbx-mcp/tests/protocol.rs index 8f81079672..4d91ad58bf 100644 --- a/crates/dbx-mcp/tests/protocol.rs +++ b/crates/dbx-mcp/tests/protocol.rs @@ -1,4 +1,4 @@ -use std::sync::Arc; +use std::{collections::HashMap, sync::Arc}; use async_trait::async_trait; use dbx_core::{ @@ -50,6 +50,7 @@ impl DbxBackend for EmptyBackend { struct PolicyBackend { policy: McpGlobalPolicy, connections: Vec, + group_paths: Result>, String>, } #[async_trait] @@ -62,6 +63,10 @@ impl DbxBackend for PolicyBackend { Ok(self.connections.clone()) } + async fn load_connection_group_paths(&self) -> Result>, String> { + self.group_paths.clone() + } + async fn execute_agent_tool( &self, _connection: &ConnectionConfig, @@ -133,9 +138,18 @@ async fn enforces_global_connection_scope_and_read_only_policy() { policy: McpGlobalPolicy { read_only: true, allow_dangerous_sql: false, - allowed_connection_ids: Some(vec!["allowed".to_string()]), + allowed_connection_ids: Some(vec!["allowed".to_string(), "allowed-staging".to_string()]), }, - connections: vec![test_connection("allowed", "allowed-db"), test_connection("blocked", "blocked-db")], + connections: vec![ + test_connection("allowed", "shared-db"), + test_connection("allowed-staging", "shared-db"), + test_connection("blocked", "blocked-db"), + ], + group_paths: Ok(HashMap::from([ + ("allowed".to_string(), vec!["Project".to_string(), "Production".to_string()]), + ("allowed-staging".to_string(), vec!["Project".to_string(), "Staging".to_string()]), + ("blocked".to_string(), vec!["Secret".to_string()]), + ])), }; let (server_transport, client_transport) = tokio::io::duplex(16 * 1024); let server = DbxMcpServer::with_runtime_options(Arc::new(backend), McpScope::default(), false); @@ -145,8 +159,11 @@ async fn enforces_global_connection_scope_and_read_only_policy() { let listed = client.peer().call_tool(CallToolRequestParams::new("dbx_list_connections")).await.expect("list connections"); let listed_text = listed.content[0].as_text().expect("list result").text.clone(); - assert!(listed_text.contains("allowed-db")); + assert_eq!(listed_text.matches("shared-db").count(), 2); assert!(!listed_text.contains("blocked-db")); + assert!(listed_text.contains("Project / Production")); + assert!(listed_text.contains("Project / Staging")); + assert!(!listed_text.contains("Secret")); let blocked = client .peer() @@ -176,3 +193,56 @@ async fn enforces_global_connection_scope_and_read_only_policy() { client.cancel().await.expect("close MCP client"); server_task.abort(); } + +#[tokio::test] +async fn connection_group_path_failure_preserves_connection_listing() { + let backend = PolicyBackend { + policy: McpGlobalPolicy::default(), + connections: vec![test_connection("local", "local-db")], + group_paths: Err("layout unavailable".to_string()), + }; + let (server_transport, client_transport) = tokio::io::duplex(16 * 1024); + let server = DbxMcpServer::with_runtime_options(Arc::new(backend), McpScope::default(), false); + let server_task = tokio::spawn(async move { server.serve(server_transport).await }); + let client = ().serve(client_transport).await.expect("initialize MCP client"); + + let listed = + client.peer().call_tool(CallToolRequestParams::new("dbx_list_connections")).await.expect("list connections"); + let listed_text = listed.content[0].as_text().expect("list result").text.clone(); + assert_ne!(listed.is_error, Some(true)); + assert!(listed_text.contains("| ID | Name | Group Path |")); + assert!(listed_text.contains("local-db")); + + client.cancel().await.expect("close MCP client"); + server_task.abort(); +} + +#[tokio::test] +async fn runtime_connection_scope_preserves_group_paths() { + let backend = PolicyBackend { + policy: McpGlobalPolicy::default(), + connections: vec![test_connection("scoped", "shared-db"), test_connection("outside", "shared-db")], + group_paths: Ok(HashMap::from([ + ("scoped".to_string(), vec!["Project".to_string(), "Production".to_string()]), + ("outside".to_string(), vec!["Project".to_string(), "Staging".to_string()]), + ])), + }; + let (server_transport, client_transport) = tokio::io::duplex(16 * 1024); + let server = DbxMcpServer::with_runtime_options( + Arc::new(backend), + McpScope { connection_ids: vec!["scoped".to_string()], ..Default::default() }, + false, + ); + let server_task = tokio::spawn(async move { server.serve(server_transport).await }); + let client = ().serve(client_transport).await.expect("initialize MCP client"); + + let listed = + client.peer().call_tool(CallToolRequestParams::new("dbx_list_connections")).await.expect("list connections"); + let listed_text = listed.content[0].as_text().expect("list result").text.clone(); + assert!(listed_text.contains("| scoped | shared-db | Project / Production |")); + assert!(!listed_text.contains("outside")); + assert!(!listed_text.contains("Project / Staging")); + + client.cancel().await.expect("close MCP client"); + server_task.abort(); +}