diff --git a/apps/desktop/src/lib/__tests__/ai/ai.spec.ts b/apps/desktop/src/lib/__tests__/ai/ai.spec.ts index e52fda91f8..304eb398d9 100644 --- a/apps/desktop/src/lib/__tests__/ai/ai.spec.ts +++ b/apps/desktop/src/lib/__tests__/ai/ai.spec.ts @@ -57,4 +57,13 @@ describe("AI SQL dialect prompt", () => { expect(prompt).not.toContain("Never execute writes without confirmation"); expect(prompt).toContain("Ask mode"); }); + + it("MongoDB agent mode uses shell commands instead of SQL", () => { + const prompt = buildSystemPrompt("general", context({ databaseType: "mongodb", connectionName: "MongoDB", database: "benchmark" }), "agent"); + + expect(prompt).toContain("MongoDB Agent mode"); + expect(prompt).toContain("execute_query accepts MongoDB shell-style commands, not SQL"); + expect(prompt).toContain("db.collection.findOne({})"); + expect(prompt).not.toContain("get_sample_data"); + }); }); diff --git a/apps/desktop/src/lib/ai/ai.ts b/apps/desktop/src/lib/ai/ai.ts index 225ac22a37..2da72ee973 100644 --- a/apps/desktop/src/lib/ai/ai.ts +++ b/apps/desktop/src/lib/ai/ai.ts @@ -253,7 +253,7 @@ export function buildSystemPrompt(action: AiAction, context: AiContext, mode: Ai const isZh = isChineseLocale(currentLocale()); - const lines: string[] = [...buildBasePromptLines(isZh), ...buildModePromptLines(mode, isZh), ...buildActionPromptLines(action, isZh), ...buildCustomInstructionLines(custom, isZh)]; + const lines: string[] = [...buildBasePromptLines(isZh), ...buildModePromptLines(mode, isZh, context.databaseType), ...buildActionPromptLines(action, isZh), ...buildCustomInstructionLines(custom, isZh)]; if (schemaScope === "focused_table") { lines.push( @@ -363,8 +363,24 @@ function buildVectorModePromptLines(context: AiContext, mode: AiAssistantMode, i ]; } -function buildModePromptLines(mode: AiAssistantMode, isZh: boolean): string[] { +function buildModePromptLines(mode: AiAssistantMode, isZh: boolean, databaseType: DatabaseType): string[] { const currentTimeGuidance = currentTimeToolGuidance(); + if (databaseType === "mongodb") { + if (mode === "agent") { + return [ + isZh ? "你处于 MongoDB Agent 模式。你有以下工具可用：list_tables、get_columns、execute_query、get_current_time。" : "You are in MongoDB Agent mode. You have the following tools available: list_tables, get_columns, execute_query, get_current_time.", + isZh + ? "execute_query 接收 MongoDB shell 风格命令，不是 SQL，例如 db.collection.find({})、db.collection.findOne({})、db.collection.aggregate([])。用户提出数据查询意图时，必须调用该工具获取真实结果后再回答。" + : "execute_query accepts MongoDB shell-style commands, not SQL, for example db.collection.find({}), db.collection.findOne({}), or db.collection.aggregate([]). For data queries, call the tool and answer from its actual results.", + currentTimeGuidance, + isZh ? "禁止不经确认直接执行 MongoDB 写命令；如果安全执行条件不满足，先说明原因，再给只读预览或澄清问题。" : "Never execute MongoDB write commands without confirmation. If safe execution requirements are not met, explain why first, then provide a read-only preview or a clarifying question.", + ]; + } + return [ + isZh ? "你处于 MongoDB Ask 模式。只生成 MongoDB shell 风格命令和说明，不要生成 SQL，也不要暗示已经执行或即将自动执行。" : "You are in MongoDB Ask mode. Generate MongoDB shell-style commands and explanations, not SQL, and do not imply that anything has run or will auto-run.", + currentTimeGuidance, + ]; + } if (mode === "agent") { return [ isZh ? "你处于 Agent 模式。你有以下工具可用：list_tables、get_columns、execute_query、get_sample_data、get_current_time。" : "You are in Agent mode. You have the following tools available: list_tables, get_columns, execute_query, get_sample_data, get_current_time.", diff --git a/crates/dbx-core/src/agent_tools.rs b/crates/dbx-core/src/agent_tools.rs index bae1fca34a..eee9f58aa4 100644 --- a/crates/dbx-core/src/agent_tools.rs +++ b/crates/dbx-core/src/agent_tools.rs @@ -191,7 +191,9 @@ pub fn all_tools(db_type: DatabaseType, sql_permissions: AgentSqlPermissions) -> return vec![list_collections_tool(), browse_collection_tool(), get_current_time_tool()]; } let mut tools = vec![list_tables_tool(), get_columns_tool(), get_current_time_tool()]; - if supports_sql_query(db_type) { + if db_type == DatabaseType::MongoDb { + tools.push(mongo_execute_query_tool(sql_permissions)); + } else if supports_sql_query(db_type) { tools.push(execute_query_tool(sql_permissions)); tools.push(get_sample_data_tool()); } @@ -201,6 +203,29 @@ pub fn all_tools(db_type: DatabaseType, sql_permissions: AgentSqlPermissions) -> tools } +fn mongo_execute_query_tool(_sql_permissions: AgentSqlPermissions) -> ToolDefinition { + ToolDefinition { + name: "execute_query", + description: "Execute a read-only MongoDB shell command and return results (max 50 rows). Use commands such as db.collection.find({}), db.collection.findOne({}), db.collection.aggregate([]), or db.collection.countDocuments({}). Write commands are not available to the MongoDB Agent.", + parameters: json!({ + "type": "object", + "properties": { + "sql": { + "type": "string", + "description": "The MongoDB shell-style command to execute" + }, + "limit": { + "type": "number", + "description": "Max rows to return (default 50, max 100)" + } + }, + "required": ["sql"] + }), + read_only: true, + parallel_ok: false, + } +} + /// list_tables tool definition. fn list_tables_tool() -> ToolDefinition { ToolDefinition { @@ -620,6 +645,10 @@ async fn execute_execute_query( .map(|l| (l as usize).min(MAX_ALLOWED_ROWS)) .unwrap_or(EXECUTE_QUERY_LIMIT); + if *db_type == DatabaseType::MongoDb { + return execute_mongo_query(state, connection_id, database, sql, limit.max(1)).await; + } + // Classify SQL risk using the concrete database dialect. let risk = crate::sql_risk::classify_sql_risk_for_database(sql, *db_type)?; let connection_config = state.configs.read().await.get(connection_id).cloned(); @@ -677,6 +706,29 @@ async fn execute_execute_query( format_query_result_as_text(&result, limit) } +async fn execute_mongo_query( + state: &Arc, + connection_id: &str, + database: &str, + source: &str, + limit: usize, +) -> Result { + let command = crate::mongo_shell::parse(source).map_err(|error| { + format!( + "{error} Use MongoDB shell-style commands such as db.collection.find({{}}), db.collection.findOne({{}}), or db.collection.aggregate([])." + ) + })?; + if command.is_mutating() { + return Err( + "Blocked: MongoDB Agent queries are read-only. Return the command for the user to review and execute manually in DBX." + .to_string(), + ); + } + + let result = crate::mongo_ops::execute_mongo_command_core(state, connection_id, database, &command, limit).await?; + format_query_result_as_text(&result, limit) +} + /// Format a QueryResult as a Markdown table for LLM consumption. fn format_query_result_as_text(result: &QueryResult, limit: usize) -> Result { // A result without columns is a command result, not an empty result set. @@ -1174,6 +1226,94 @@ for line in sys.stdin: assert!(names.contains(&"get_current_time")); } + #[cfg(unix)] + #[tokio::test] + async fn mongodb_agent_registers_shell_query_tool_and_routes_find_one_as_read_only() { + let tools = all_tools(DatabaseType::MongoDb, AgentSqlPermissions::default()); + let names = tools.iter().map(|tool| tool.name).collect::>(); + assert!(names.contains(&"execute_query")); + assert!(!names.contains(&"get_sample_data")); + assert!(!names.contains(&"explain_query")); + let execute_query = tools.iter().find(|tool| tool.name == "execute_query").unwrap(); + assert!(execute_query.description.contains("MongoDB shell command")); + assert!(!execute_query.description.contains("SQL query")); + assert!(execute_query.description.contains("read-only")); + + let confirmed_tools = all_tools( + DatabaseType::MongoDb, + AgentSqlPermissions { + allow_writes: true, + allow_dangerous: true, + confirmed_write_sql: Some("db.items.insertOne({name: 'test'})".to_string()), + }, + ); + let confirmed_execute_query = confirmed_tools.iter().find(|tool| tool.name == "execute_query").unwrap(); + assert!(confirmed_execute_query.description.contains("read-only")); + assert!(!confirmed_execute_query.description.contains("confirmed write")); + + let temp_dir = tempfile::tempdir().unwrap(); + let storage = Storage::open(&temp_dir.path().join("storage.db")).await.unwrap(); + let state = Arc::new(AppState::new(storage)); + let call = ToolCall { + id: "mongo-find-one".to_string(), + name: "execute_query".to_string(), + arguments: json!({ "sql": "db.BenchmarkIndex_approved.findOne({})" }), + provider_payload: None, + }; + + let result = execute_tool( + &call, + &state, + "mongo-1", + "benchmark", + None, + &DatabaseType::MongoDb, + AgentSqlPermissions::default(), + ) + .await; + + assert!(result.is_error); + assert!(!result.content.contains("Blocked:"), "{}", result.content); + assert!(result.content.contains("Connection") || result.content.contains("connection"), "{}", result.content); + } + + #[cfg(unix)] + #[tokio::test] + async fn mongodb_agent_keeps_all_writes_blocked_after_sql_confirmation() { + let temp_dir = tempfile::tempdir().unwrap(); + let storage = Storage::open(&temp_dir.path().join("storage.db")).await.unwrap(); + let state = Arc::new(AppState::new(storage)); + let permissions = AgentSqlPermissions { + allow_writes: true, + allow_dangerous: true, + confirmed_write_sql: Some("SQL confirmation does not grant MongoDB writes".to_string()), + }; + + for (index, source) in [ + "db.items.insertOne({name: 'test'})", + "db.items.updateMany({tenant: 7}, {$set: {active: false}})", + "db.items.deleteMany({tenant: 7})", + "db.items.createIndex({tenant: 1})", + r#"db.items.aggregate([{"$out":"items_backup"}])"#, + ] + .into_iter() + .enumerate() + { + let call = ToolCall { + id: format!("mongo-write-{index}"), + name: "execute_query".to_string(), + arguments: json!({ "sql": source }), + provider_payload: None, + }; + let result = + execute_tool(&call, &state, "mongo-1", "benchmark", None, &DatabaseType::MongoDb, permissions.clone()) + .await; + + assert!(result.is_error, "{source}: {}", result.content); + assert!(result.content.contains("read-only"), "{source}: {}", result.content); + } + } + #[test] fn confirmed_sql_permissions_update_execute_query_contract() { let tools = all_tools( diff --git a/crates/dbx-core/src/db/mongo_driver.rs b/crates/dbx-core/src/db/mongo_driver.rs index 2aa7299440..583c68cb5b 100644 --- a/crates/dbx-core/src/db/mongo_driver.rs +++ b/crates/dbx-core/src/db/mongo_driver.rs @@ -709,6 +709,41 @@ pub async fn find_documents( projection: Option<&str>, sort: Option<&str>, collation: Option<&str>, +) -> Result { + find_documents_with_total(client, database, collection, skip, limit, filter, projection, sort, collation, true) + .await +} + +/// Execute a find without the document browser's separate total-count query. +/// Agent callers only consume returned rows, so counting the full result set adds latency +/// without providing any useful output. +pub async fn find_documents_without_total( + client: &Client, + database: &str, + collection: &str, + skip: u64, + limit: i64, + filter: Option<&str>, + projection: Option<&str>, + sort: Option<&str>, + collation: Option<&str>, +) -> Result { + find_documents_with_total(client, database, collection, skip, limit, filter, projection, sort, collation, false) + .await +} + +#[allow(clippy::too_many_arguments)] +async fn find_documents_with_total( + client: &Client, + database: &str, + collection: &str, + skip: u64, + limit: i64, + filter: Option<&str>, + projection: Option<&str>, + sort: Option<&str>, + collation: Option<&str>, + include_total: bool, ) -> Result { let col = client.database(database).collection::(collection); @@ -722,14 +757,16 @@ pub async fn find_documents( let collation = parse_find_collation(collation)?; let count_is_exact = !filter_doc.is_empty(); - let total_result = if count_is_exact { + let total_result = if !include_total { + None + } else if count_is_exact { let mut count = col.count_documents(filter_doc.clone()); if let Some(collation) = collation.clone() { count = count.collation(collation); } - count.await.map_err(|e| e.to_string()) + Some(count.await.map_err(|e| e.to_string())) } else { - col.estimated_document_count().await.map_err(|e| e.to_string()) + Some(col.estimated_document_count().await.map_err(|e| e.to_string())) }; let mut find = col.find(filter_doc).skip(skip).limit(limit); @@ -761,7 +798,10 @@ pub async fn find_documents( documents.push(bson_to_json(&Bson::Document(doc.clone()))); extended_documents.push(Bson::Document(doc).into_canonical_extjson()); } - let (total, total_is_exact) = resolve_mongo_find_total(total_result, count_is_exact, skip, documents.len()); + let (total, total_is_exact) = match total_result { + Some(total_result) => resolve_mongo_find_total(total_result, count_is_exact, skip, documents.len()), + None => (documents.len() as u64, false), + }; Ok(MongoDocumentResult { documents, diff --git a/crates/dbx-core/src/mongo_ops.rs b/crates/dbx-core/src/mongo_ops.rs index 69f2a61f25..0d5bb902ed 100644 --- a/crates/dbx-core/src/mongo_ops.rs +++ b/crates/dbx-core/src/mongo_ops.rs @@ -4,6 +4,8 @@ use crate::db::mongo_driver::{ self, MongoCollectionStatsResult, MongoDocumentResult, MongoDropIndexFailure, MongoDropIndexesResult, }; use crate::document_ops::CollectionInfo; +use crate::mongo_shell::MongoCommand; +use crate::types::QueryResult; async fn ensure_document_pool(state: &AppState, connection_id: &str) -> Result<(), String> { state.get_or_create_pool(connection_id, None).await.map(|_| ()) @@ -155,6 +157,50 @@ pub async fn mongo_find_documents_core( .await } +#[allow(clippy::too_many_arguments)] +async fn mongo_find_documents_without_total_core( + state: &AppState, + connection_id: &str, + database: &str, + collection: &str, + skip: u64, + limit: i64, + filter: Option<&str>, + projection: Option<&str>, + sort: Option<&str>, + collation: Option<&str>, +) -> Result { + ensure_document_pool(state, connection_id).await?; + let connections = state.connections.read().await; + match connections.get(connection_id).ok_or("Not found")? { + PoolKind::MongoDb(client) => { + mongo_driver::find_documents_without_total( + client, database, collection, skip, limit, filter, projection, sort, collation, + ) + .await + } + PoolKind::Agent(client) => { + let mut client = client.lock().await; + let mut params = serde_json::json!({ + "database": database, + "collection": collection, + "skip": skip, + "limit": limit, + "filter": filter, + "sort": sort, + }); + if let Some(projection) = projection { + params["projection"] = serde_json::json!(projection); + } + if let Some(collation) = collation { + params["collation"] = serde_json::json!(collation); + } + client.mongo_find_documents(params).await + } + _ => Err("Not a MongoDB connection".to_string()), + } +} + pub async fn mongo_find_one_core( state: &AppState, connection_id: &str, @@ -635,10 +681,311 @@ pub async fn mongo_find_one_and_delete_core( } } +pub async fn execute_mongo_command_core( + state: &AppState, + connection_id: &str, + database: &str, + command: &MongoCommand, + max_rows: usize, +) -> Result { + use serde_json::Value; + + match command { + MongoCommand::Version => mongo_server_version_core(state, connection_id, database) + .await + .map(|version| scalar_query_result("version", Value::String(version))), + MongoCommand::Use { database } => Ok(scalar_query_result("database", Value::String(database.clone()))), + MongoCommand::Find { collection, filter, projection, sort, collation, skip, limit } => { + let limit = bounded_mongo_find_limit(*limit, max_rows); + let result = mongo_find_documents_without_total_core( + state, + connection_id, + database, + collection, + *skip, + limit, + Some(filter), + projection.as_deref(), + sort.as_deref(), + collation.as_deref(), + ) + .await?; + Ok(mongo_documents_query_result(result.documents)) + } + MongoCommand::FindOne { collection, filter, projection, options } => { + let result = mongo_find_one_core( + state, + connection_id, + database, + collection, + Some(filter), + projection.as_deref(), + options.as_deref(), + ) + .await?; + Ok(mongo_documents_query_result(result.documents)) + } + MongoCommand::Count { collection, filter, accurate } => { + let mode = if *accurate { "accurate" } else { "legacy" }; + let total = + mongo_count_documents_core(state, connection_id, database, collection, Some(filter), Some(mode)) + .await?; + Ok(scalar_query_result("count", Value::from(total))) + } + MongoCommand::Aggregate { collection, pipeline, options } => { + let result = mongo_aggregate_documents_core( + state, + connection_id, + database, + collection, + pipeline, + Some(max_rows), + options.as_deref(), + ) + .await?; + Ok(mongo_documents_query_result(result.documents)) + } + MongoCommand::Distinct { collection, field, filter } => { + let result = + mongo_distinct_core(state, connection_id, database, collection, field, filter.as_deref()).await?; + Ok(mongo_documents_query_result(limit_mongo_documents(result, max_rows).documents)) + } + MongoCommand::GetIndexes { collection } => { + let result = mongo_aggregate_documents_core( + state, + connection_id, + database, + collection, + r#"[{"$indexStats":{}}]"#, + Some(max_rows), + None, + ) + .await?; + Ok(mongo_documents_query_result(result.documents)) + } + MongoCommand::CollectionStats { collection, metric, scale } => { + let stats = mongo_collection_stats_core(state, connection_id, database, collection, scale.clone()).await?; + let value = serde_json::to_value(stats).map_err(|error| error.to_string())?; + if metric == "stats" { + Ok(mongo_documents_query_result(vec![value])) + } else { + let key = match metric.as_str() { + "dataSize" => "size", + "storageSize" => "storageSize", + "totalIndexSize" => "totalIndexSize", + _ => metric, + }; + Ok(scalar_query_result(metric, value.get(key).cloned().unwrap_or(Value::Null))) + } + } + MongoCommand::Insert { collection, documents } => { + let affected = mongo_insert_documents_core(state, connection_id, database, collection, documents).await?; + Ok(affected_query_result(affected)) + } + MongoCommand::Update { collection, filter, update, options, many } => { + let affected = mongo_update_documents_core( + state, + connection_id, + database, + collection, + filter, + update, + *many, + options.as_deref(), + ) + .await?; + Ok(affected_query_result(affected)) + } + MongoCommand::Delete { collection, filter, many } => { + let affected = + mongo_delete_documents_core(state, connection_id, database, collection, filter, *many).await?; + Ok(affected_query_result(affected)) + } + MongoCommand::CreateIndex { collection, keys, options } => { + let name = + mongo_create_index_core(state, connection_id, database, collection, keys, options.as_deref()).await?; + Ok(scalar_query_result("name", Value::String(name))) + } + MongoCommand::DropIndexes { collection, indexes, single } => { + let result = + mongo_drop_indexes_core(state, connection_id, database, collection, indexes.as_deref(), *single) + .await?; + Ok(mongo_drop_indexes_query_result( + result.dropped_names, + result.failures.into_iter().map(|failure| (failure.name, failure.message)).collect(), + result.affected_rows, + )) + } + MongoCommand::DropCollection { collection } => { + mongo_drop_collection_core(state, connection_id, database, collection).await?; + Ok(affected_query_result(1)) + } + MongoCommand::FindOneAndUpdate { collection, filter, update, options } => { + let result = mongo_find_one_and_update_core( + state, + connection_id, + database, + collection, + filter, + update, + options.as_deref(), + ) + .await?; + Ok(mongo_documents_query_result(result.documents)) + } + MongoCommand::FindOneAndReplace { collection, filter, replacement, options } => { + let result = mongo_find_one_and_replace_core( + state, + connection_id, + database, + collection, + filter, + replacement, + options.as_deref(), + ) + .await?; + Ok(mongo_documents_query_result(result.documents)) + } + MongoCommand::FindOneAndDelete { collection, filter, options } => { + let result = + mongo_find_one_and_delete_core(state, connection_id, database, collection, filter, options.as_deref()) + .await?; + Ok(mongo_documents_query_result(result.documents)) + } + } +} + +fn bounded_mongo_find_limit(command_limit: i64, max_rows: usize) -> i64 { + let max_rows = max_rows.max(1).min(i64::MAX as usize) as i64; + if command_limit == 0 { + return max_rows; + } + command_limit.saturating_abs().min(max_rows).max(1) +} + +fn limit_mongo_documents(mut result: MongoDocumentResult, max_rows: usize) -> MongoDocumentResult { + let max_rows = max_rows.max(1); + result.documents.truncate(max_rows); + if let Some(raw_documents) = result.raw_documents.as_mut() { + raw_documents.truncate(max_rows); + } + if let Some(extended_documents) = result.extended_documents.as_mut() { + extended_documents.truncate(max_rows); + } + result +} + +fn query_result(columns: Vec, rows: Vec>, affected_rows: u64) -> QueryResult { + QueryResult { + columns, + column_types: Vec::new(), + column_sortables: Vec::new(), + spatial_columns: Vec::new(), + spatial_values: Vec::new(), + rows, + affected_rows, + execution_time_ms: 0, + truncated: false, + session_id: None, + has_more: false, + elasticsearch_raw_body: None, + } +} + +fn scalar_query_result(column: impl Into, value: serde_json::Value) -> QueryResult { + query_result(vec![column.into()], vec![vec![value]], 0) +} + +fn affected_query_result(affected_rows: u64) -> QueryResult { + query_result(Vec::new(), Vec::new(), affected_rows) +} + +fn mongo_drop_indexes_query_result( + dropped_names: Vec, + failures: Vec<(String, String)>, + affected_rows: u64, +) -> QueryResult { + use serde_json::Value; + + if failures.is_empty() { + let rows = dropped_names.into_iter().map(|name| vec![Value::String(name)]).collect::>(); + return query_result(if rows.is_empty() { Vec::new() } else { vec!["name".to_string()] }, rows, affected_rows); + } + + let mut rows = dropped_names + .into_iter() + .map(|name| vec![Value::String(name), Value::String("dropped".to_string()), Value::Null]) + .collect::>(); + rows.extend( + failures.into_iter().map(|(name, message)| { + vec![Value::String(name), Value::String("failed".to_string()), Value::String(message)] + }), + ); + query_result(vec!["name".to_string(), "status".to_string(), "message".to_string()], rows, affected_rows) +} + +fn mongo_documents_query_result(documents: Vec) -> QueryResult { + use serde_json::Value; + + if documents.is_empty() { + return query_result(Vec::new(), Vec::new(), 0); + } + let mut columns = std::collections::BTreeSet::new(); + for document in &documents { + if let Some(object) = document.as_object() { + columns.extend(object.keys().cloned()); + } else { + columns.insert("value".to_string()); + } + } + let columns = columns.into_iter().collect::>(); + let rows = documents + .into_iter() + .map(|document| { + columns + .iter() + .map(|column| { + document + .as_object() + .and_then(|object| object.get(column)) + .cloned() + .or_else(|| (column == "value").then(|| document.clone())) + .unwrap_or(Value::Null) + }) + .collect() + }) + .collect(); + query_result(columns, rows, 0) +} + #[cfg(test)] mod tests { use super::*; + #[test] + fn agent_find_limit_never_passes_an_unbounded_zero_to_mongodb() { + assert_eq!(bounded_mongo_find_limit(0, 50), 50); + assert_eq!(bounded_mongo_find_limit(0, 0), 1); + assert_eq!(bounded_mongo_find_limit(100, 7), 7); + assert_eq!(bounded_mongo_find_limit(-100, 7), 7); + } + + #[test] + fn agent_document_limit_caps_distinct_values() { + let result = MongoDocumentResult { + documents: vec![serde_json::json!(1), serde_json::json!(2), serde_json::json!(3)], + raw_documents: Some(vec!["1".to_string(), "2".to_string(), "3".to_string()]), + extended_documents: Some(vec![serde_json::json!(1), serde_json::json!(2), serde_json::json!(3)]), + total: 3, + total_is_exact: true, + }; + + let limited = limit_mongo_documents(result, 2); + assert_eq!(limited.documents, vec![serde_json::json!(1), serde_json::json!(2)]); + assert_eq!(limited.raw_documents.unwrap(), vec!["1", "2"]); + assert_eq!(limited.extended_documents.unwrap(), vec![serde_json::json!(1), serde_json::json!(2)]); + } + #[cfg(unix)] use crate::db::agent_driver::{AgentDriverClient, AgentLaunchSpec}; #[cfg(unix)] diff --git a/crates/dbx-core/tests/live_mongodb_agent_tools.rs b/crates/dbx-core/tests/live_mongodb_agent_tools.rs new file mode 100644 index 0000000000..2d8caf5f20 --- /dev/null +++ b/crates/dbx-core/tests/live_mongodb_agent_tools.rs @@ -0,0 +1,130 @@ +use std::sync::{Arc, Mutex}; + +use dbx_core::agent_events::{ToolCall, ToolResult}; +use dbx_core::agent_tools::{execute_tool, AgentSqlPermissions}; +use dbx_core::connection::{AppState, PoolKind}; +use dbx_core::models::connection::{ConnectionConfig, DatabaseType}; +use dbx_core::storage::Storage; +use mongodb::event::{command::CommandEvent, EventHandler}; +use mongodb::options::ClientOptions; +use mongodb::Client; + +fn live_mongodb_config(id: &str) -> ConnectionConfig { + let host = std::env::var("DBX_LIVE_MONGODB_HOST").expect("DBX_LIVE_MONGODB_HOST"); + let port = std::env::var("DBX_LIVE_MONGODB_PORT").ok().and_then(|value| value.parse().ok()).unwrap_or(27017); + let username = std::env::var("DBX_LIVE_MONGODB_USER").expect("DBX_LIVE_MONGODB_USER"); + let password = std::env::var("DBX_LIVE_MONGODB_PASSWORD").expect("DBX_LIVE_MONGODB_PASSWORD"); + let database = std::env::var("DBX_LIVE_MONGODB_DATABASE").expect("DBX_LIVE_MONGODB_DATABASE"); + let auth_source = std::env::var("DBX_LIVE_MONGODB_AUTH_SOURCE").unwrap_or_else(|_| database.clone()); + + serde_json::from_value(serde_json::json!({ + "id": id, + "name": "Live MongoDB agent tools", + "db_type": DatabaseType::MongoDb, + "host": host, + "port": port, + "username": username, + "password": password, + "database": database, + "url_params": format!("authSource={auth_source}"), + "connect_timeout_secs": 10, + "query_timeout_secs": 30, + "idle_timeout_secs": 60, + "keepalive_interval_secs": 0 + })) + .expect("live MongoDB connection config should deserialize") +} + +async fn call_agent_tool(state: &Arc, connection_id: &str, database: &str, source: &str) -> ToolResult { + call_agent_tool_with_limit(state, connection_id, database, source, None).await +} + +async fn call_agent_tool_with_limit( + state: &Arc, + connection_id: &str, + database: &str, + source: &str, + limit: Option, +) -> ToolResult { + let mut arguments = serde_json::json!({ "sql": source }); + if let Some(limit) = limit { + arguments["limit"] = serde_json::json!(limit); + } + execute_tool( + &ToolCall { + id: "live-mongodb-agent".to_string(), + name: "execute_query".to_string(), + arguments, + provider_payload: None, + }, + state, + connection_id, + database, + None, + &DatabaseType::MongoDb, + AgentSqlPermissions::default(), + ) + .await +} + +#[tokio::test] +#[ignore = "requires DBX_LIVE_MONGODB_* env vars pointing at a MongoDB database with a non-empty stores collection"] +async fn mongodb_agent_find_one_returns_real_data_and_keeps_writes_blocked() { + let directory = tempfile::tempdir().unwrap(); + let storage = Storage::open(&directory.path().join("storage.db")).await.unwrap(); + let state = Arc::new(AppState::new(storage)); + let config = live_mongodb_config("live-mongodb-agent"); + let database = config.database.clone().expect("database"); + state.configs.write().await.insert(config.id.clone(), config.clone()); + + let read = call_agent_tool(&state, &config.id, &database, "db.stores.findOne({})").await; + assert!(!read.is_error, "{}", read.content); + assert!(read.content.contains("(1 rows,"), "{}", read.content); + + let write = + call_agent_tool(&state, &config.id, &database, "db.agent_regression.insertOne({marker: 'blocked'})").await; + assert!(write.is_error); + assert!(write.content.contains("read-only"), "{}", write.content); +} + +#[tokio::test] +#[ignore = "requires DBX_LIVE_MONGODB_* env vars pointing at a MongoDB database with a products collection containing at least 8 documents"] +async fn mongodb_agent_enforces_limits_and_find_skips_total_count() { + let directory = tempfile::tempdir().unwrap(); + let storage = Storage::open(&directory.path().join("storage.db")).await.unwrap(); + let state = Arc::new(AppState::new(storage)); + let config = live_mongodb_config("live-mongodb-agent-limits"); + let database = config.database.clone().expect("database"); + + let command_names = Arc::new(Mutex::new(Vec::::new())); + let observed_names = command_names.clone(); + let mut options = ClientOptions::parse(config.connection_url()).await.unwrap(); + options.command_event_handler = Some(EventHandler::callback(move |event| { + if let CommandEvent::Started(event) = event { + observed_names.lock().unwrap().push(event.command_name); + } + })); + let client = Client::with_options(options).unwrap(); + + state.configs.write().await.insert(config.id.clone(), config.clone()); + state.connections.write().await.insert(config.id.clone(), PoolKind::MongoDb(client)); + + let zero_limit = + call_agent_tool_with_limit(&state, &config.id, &database, "db.products.find({}).limit(0)", Some(0)).await; + assert!(!zero_limit.is_error, "{}", zero_limit.content); + assert!(zero_limit.content.contains("(1 rows,"), "{}", zero_limit.content); + + let commands = command_names.lock().unwrap().clone(); + assert!(commands.iter().any(|command| command == "find"), "{commands:?}"); + assert!(!commands.iter().any(|command| command == "count" || command == "aggregate"), "{commands:?}"); + + let bounded = + call_agent_tool_with_limit(&state, &config.id, &database, "db.products.find({}).limit(0)", Some(7)).await; + assert!(!bounded.is_error, "{}", bounded.content); + assert!(bounded.content.contains("(7 rows,"), "{}", bounded.content); + + let distinct = + call_agent_tool_with_limit(&state, &config.id, &database, "db.products.distinct('_id')", Some(7)).await; + assert!(!distinct.is_error, "{}", distinct.content); + assert!(distinct.content.contains("(7 rows,"), "{}", distinct.content); +} diff --git a/crates/dbx-mcp/src/backend.rs b/crates/dbx-mcp/src/backend.rs index ed5d91193a..75c850a6db 100644 --- a/crates/dbx-mcp/src/backend.rs +++ b/crates/dbx-mcp/src/backend.rs @@ -461,220 +461,7 @@ impl DbxBackend for LocalBackend { database: &str, command: &MongoCommand, ) -> Result { - use dbx_core::mongo_ops; - - let connection_id = &connection.id; - match command { - MongoCommand::Version => mongo_ops::mongo_server_version_core(&self.state, connection_id, database) - .await - .map(|version| scalar_query_result("version", Value::String(version))), - MongoCommand::Use { database } => Ok(scalar_query_result("database", Value::String(database.clone()))), - MongoCommand::Find { collection, filter, projection, sort, collation, skip, limit } => { - let result = mongo_ops::mongo_find_documents_core( - &self.state, - connection_id, - database, - collection, - *skip, - *limit, - Some(filter), - projection.as_deref(), - sort.as_deref(), - collation.as_deref(), - ) - .await?; - Ok(mongo_documents_query_result(result.documents)) - } - MongoCommand::FindOne { collection, filter, projection, options } => { - let result = mongo_ops::mongo_find_one_core( - &self.state, - connection_id, - database, - collection, - Some(filter), - projection.as_deref(), - options.as_deref(), - ) - .await?; - Ok(mongo_documents_query_result(result.documents)) - } - MongoCommand::Count { collection, filter, accurate } => { - let mode = if *accurate { "accurate" } else { "legacy" }; - let total = mongo_ops::mongo_count_documents_core( - &self.state, - connection_id, - database, - collection, - Some(filter), - Some(mode), - ) - .await?; - Ok(scalar_query_result("count", Value::from(total))) - } - MongoCommand::Aggregate { collection, pipeline, options } => { - let result = mongo_ops::mongo_aggregate_documents_core( - &self.state, - connection_id, - database, - collection, - pipeline, - Some(100), - options.as_deref(), - ) - .await?; - Ok(mongo_documents_query_result(result.documents)) - } - MongoCommand::Distinct { collection, field, filter } => { - let result = mongo_ops::mongo_distinct_core( - &self.state, - connection_id, - database, - collection, - field, - filter.as_deref(), - ) - .await?; - Ok(mongo_documents_query_result(result.documents)) - } - MongoCommand::GetIndexes { collection } => { - let result = mongo_ops::mongo_aggregate_documents_core( - &self.state, - connection_id, - database, - collection, - r#"[{"$indexStats":{}}]"#, - Some(100), - None, - ) - .await?; - Ok(mongo_documents_query_result(result.documents)) - } - MongoCommand::CollectionStats { collection, metric, scale } => { - let stats = mongo_ops::mongo_collection_stats_core( - &self.state, - connection_id, - database, - collection, - scale.clone(), - ) - .await?; - let value = serde_json::to_value(stats).map_err(|error| error.to_string())?; - if metric == "stats" { - Ok(mongo_documents_query_result(vec![value])) - } else { - let key = match metric.as_str() { - "dataSize" => "size", - "storageSize" => "storageSize", - "totalIndexSize" => "totalIndexSize", - _ => metric, - }; - let metric_value = value.get(key).cloned().unwrap_or(Value::Null); - Ok(scalar_query_result(metric, metric_value)) - } - } - MongoCommand::Insert { collection, documents } => { - let affected = - mongo_ops::mongo_insert_documents_core(&self.state, connection_id, database, collection, documents) - .await?; - Ok(affected_query_result(affected)) - } - MongoCommand::Update { collection, filter, update, options, many } => { - let affected = mongo_ops::mongo_update_documents_core( - &self.state, - connection_id, - database, - collection, - filter, - update, - *many, - options.as_deref(), - ) - .await?; - Ok(affected_query_result(affected)) - } - MongoCommand::Delete { collection, filter, many } => { - let affected = mongo_ops::mongo_delete_documents_core( - &self.state, - connection_id, - database, - collection, - filter, - *many, - ) - .await?; - Ok(affected_query_result(affected)) - } - MongoCommand::CreateIndex { collection, keys, options } => { - let name = mongo_ops::mongo_create_index_core( - &self.state, - connection_id, - database, - collection, - keys, - options.as_deref(), - ) - .await?; - Ok(scalar_query_result("name", Value::String(name))) - } - MongoCommand::DropIndexes { collection, indexes, single } => { - let result = mongo_ops::mongo_drop_indexes_core( - &self.state, - connection_id, - database, - collection, - indexes.as_deref(), - *single, - ) - .await?; - Ok(mongo_drop_indexes_query_result( - result.dropped_names, - result.failures.into_iter().map(|failure| (failure.name, failure.message)).collect(), - result.affected_rows, - )) - } - MongoCommand::DropCollection { collection } => { - mongo_ops::mongo_drop_collection_core(&self.state, connection_id, database, collection).await?; - Ok(affected_query_result(1)) - } - MongoCommand::FindOneAndUpdate { collection, filter, update, options } => { - let result = mongo_ops::mongo_find_one_and_update_core( - &self.state, - connection_id, - database, - collection, - filter, - update, - options.as_deref(), - ) - .await?; - Ok(mongo_documents_query_result(result.documents)) - } - MongoCommand::FindOneAndReplace { collection, filter, replacement, options } => { - let result = mongo_ops::mongo_find_one_and_replace_core( - &self.state, - connection_id, - database, - collection, - filter, - replacement, - options.as_deref(), - ) - .await?; - Ok(mongo_documents_query_result(result.documents)) - } - MongoCommand::FindOneAndDelete { collection, filter, options } => { - let result = mongo_ops::mongo_find_one_and_delete_core( - &self.state, - connection_id, - database, - collection, - filter, - options.as_deref(), - ) - .await?; - Ok(mongo_documents_query_result(result.documents)) - } - } + dbx_core::mongo_ops::execute_mongo_command_core(&self.state, &connection.id, database, command, 100).await } async fn close_client_session(