From e6f8b2187a253b4bacf6f0f9e49056e2e4aeddf7 Mon Sep 17 00:00:00 2001 From: Abeautifulsnow Date: Sun, 30 Aug 2026 12:26:54 +0800 Subject: [PATCH] feat(mcp): add global query timeout override --- .../editor/EditorSettingsDialog.vue | 36 +++- ...itorSettingsDialog.mcpQueryTimeout.spec.ts | 36 ++++ apps/desktop/src/i18n/locales/en.ts | 3 + apps/desktop/src/i18n/locales/es.ts | 3 + apps/desktop/src/i18n/locales/it.ts | 3 + apps/desktop/src/i18n/locales/ja.ts | 3 + apps/desktop/src/i18n/locales/ko.ts | 3 + apps/desktop/src/i18n/locales/pt-BR.ts | 3 + apps/desktop/src/i18n/locales/zh-CN.ts | 3 + apps/desktop/src/i18n/locales/zh-TW.ts | 3 + apps/desktop/src/lib/backend/tauri.ts | 1 + .../stores/__tests__/settingsStore.spec.ts | 11 ++ apps/desktop/src/stores/settingsStore.ts | 9 + crates/dbx-core/src/agent_tools.rs | 97 +++++++++- crates/dbx-core/src/models/connection.rs | 20 +-- crates/dbx-core/src/storage.rs | 15 ++ crates/dbx-mcp/src/backend.rs | 168 +++++++++++++++++- crates/dbx-mcp/src/server.rs | 73 ++++++-- crates/dbx-mcp/tests/local.rs | 2 + crates/dbx-mcp/tests/protocol.rs | 134 +++++++++++++- crates/dbx-web/src/routes/connection.rs | 3 + crates/dbx-web/src/routes/mongo.rs | 4 + crates/dbx-web/src/routes/mq.rs | 1 + src-tauri/src/commands/mcp_bridge.rs | 16 +- 24 files changed, 615 insertions(+), 35 deletions(-) create mode 100644 apps/desktop/src/components/editor/__tests__/EditorSettingsDialog.mcpQueryTimeout.spec.ts diff --git a/apps/desktop/src/components/editor/EditorSettingsDialog.vue b/apps/desktop/src/components/editor/EditorSettingsDialog.vue index 51293607eb..a07fa7dc2d 100644 --- a/apps/desktop/src/components/editor/EditorSettingsDialog.vue +++ b/apps/desktop/src/components/editor/EditorSettingsDialog.vue @@ -2053,6 +2053,31 @@ const mcpInstallError = ref(false); const mcpExecutionMode = computed(() => mcpExecutionModeFromPolicy(settingsStore.mcpGlobalPolicy)); const mcpExecutionModeOptions: McpExecutionMode[] = ["read_only", "safe_write", "high_risk_write"]; const mcpAllowedConnectionIds = computed(() => settingsStore.mcpGlobalPolicy.allowedConnectionIds); +const mcpQueryTimeoutInput = ref(settingsStore.mcpGlobalPolicy.queryTimeoutSecs === null ? "" : String(settingsStore.mcpGlobalPolicy.queryTimeoutSecs)); + +watch( + () => settingsStore.mcpGlobalPolicy.queryTimeoutSecs, + (value) => { + mcpQueryTimeoutInput.value = value === null ? "" : String(value); + }, +); + +function onMcpQueryTimeoutInput() { + const raw = mcpQueryTimeoutInput.value.trim(); + if (raw === "") { + void saveMcpPolicy({ queryTimeoutSecs: null }); + return; + } + const parsed = Number(raw); + // Backend stores the value as u64; keep out-of-range integers on the + // invalid-value path instead of failing serde later with a generic error. + if (!Number.isFinite(parsed) || parsed < 0 || !Number.isInteger(parsed) || parsed > Number("18446744073709551615")) { + toast(t("settings.mcpQueryTimeoutInvalid"), 5000); + mcpQueryTimeoutInput.value = settingsStore.mcpGlobalPolicy.queryTimeoutSecs === null ? "" : String(settingsStore.mcpGlobalPolicy.queryTimeoutSecs); + return; + } + void saveMcpPolicy({ queryTimeoutSecs: parsed }); +} const mcpSelectableConnections = computed(() => connectionStore.connections); const mcpPolicyControlsDisabled = computed(() => isMcpPolicyMutationBlocked({ @@ -2062,7 +2087,7 @@ const mcpPolicyControlsDisabled = computed(() => }), ); -async function saveMcpPolicy(partial: { readOnly?: boolean; allowDangerousSql?: boolean; allowedConnectionIds?: string[] | null }) { +async function saveMcpPolicy(partial: { readOnly?: boolean; allowDangerousSql?: boolean; allowedConnectionIds?: string[] | null; queryTimeoutSecs?: number | null }) { if (mcpPolicyControlsDisabled.value) return; mcpPolicySaving.value = true; try { @@ -7189,6 +7214,15 @@ onUnmounted(() => { {{ t("settings.mcpExecutionModeHighRiskWriteDescription") }}

+
+
+ +

+ {{ t("settings.mcpQueryTimeoutDescription") }} +

+
+ +

diff --git a/apps/desktop/src/components/editor/__tests__/EditorSettingsDialog.mcpQueryTimeout.spec.ts b/apps/desktop/src/components/editor/__tests__/EditorSettingsDialog.mcpQueryTimeout.spec.ts new file mode 100644 index 0000000000..b257367ba1 --- /dev/null +++ b/apps/desktop/src/components/editor/__tests__/EditorSettingsDialog.mcpQueryTimeout.spec.ts @@ -0,0 +1,36 @@ +import { readFileSync } from "node:fs"; +import { describe, expect, it } from "vitest"; + +const dialogSource = readFileSync(new URL("../EditorSettingsDialog.vue", import.meta.url), "utf8"); + +// Regression for the MCP query-timeout input: decimal values like 1.5 were +// silently Math.round-ed to 2 and saved, contradicting the "non-negative whole +// number" label. The UI must reject non-integers with the existing invalid +// toast instead of silently rewriting them. +describe("EditorSettingsDialog MCP query timeout validation", () => { + function handlerBlock(): string { + const start = dialogSource.indexOf("function onMcpQueryTimeoutInput()"); + const end = dialogSource.indexOf("const mcpSelectableConnections", start); + if (start < 0 || end < 0) throw new Error("Missing onMcpQueryTimeoutInput handler block"); + return dialogSource.slice(start, end); + } + + it("rejects non-integers instead of rounding them", () => { + const block = handlerBlock(); + // The invalid path must test for integer-ness, not just finite + non-negative. + expect(block).toMatch(/Number\.isInteger\(parsed\)/); + // The valid save path must NOT wrap the value in Math.round anymore. + expect(block).not.toContain("Math.round(parsed)"); + // A decimal input falls through to the invalid toast + revert branch. + expect(block).toContain('toast(t("settings.mcpQueryTimeoutInvalid"), 5000)'); + expect(block).toMatch(/mcpQueryTimeoutInput\.value\s*=\s*settingsStore\.mcpGlobalPolicy\.queryTimeoutSecs === null \? "" : String\(settingsStore\.mcpGlobalPolicy\.queryTimeoutSecs\)/); + }); + + it("keeps null (inherit) and zero (no limit) working", () => { + const block = handlerBlock(); + // Empty input still persists null (inherit the connection). + expect(block).toContain("void saveMcpPolicy({ queryTimeoutSecs: null })"); + // A valid integer still reaches saveMcpPolicy with the exact value. + expect(block).toContain("void saveMcpPolicy({ queryTimeoutSecs: parsed })"); + }); +}); diff --git a/apps/desktop/src/i18n/locales/en.ts b/apps/desktop/src/i18n/locales/en.ts index d58d07cdb9..bc8b695df5 100644 --- a/apps/desktop/src/i18n/locales/en.ts +++ b/apps/desktop/src/i18n/locales/en.ts @@ -6911,6 +6911,9 @@ export default { mcpCapabilityConnectionManagement: "Add/remove MCP connections", mcpCapabilityAllowed: "Allowed", mcpCapabilityBlocked: "Blocked", + mcpQueryTimeout: "Query timeout (seconds)", + mcpQueryTimeoutDescription: "Leave empty to inherit the connection's query timeout. 0 = no limit. A positive value overrides the connection setting for every MCP query.", + mcpQueryTimeoutInvalid: "Enter a non-negative whole number of seconds, or leave empty to inherit the connection setting.", mcpCapabilityAlwaysEnforced: "Every mode remains subject to the connection allowlist, connection read-only protection, production protection, and database-account permissions.", mcpPolicySaveFailed: "Failed to save the MCP access policy: {error}", mcpPolicyLoadFailed: "Failed to load the MCP access policy: {error}", diff --git a/apps/desktop/src/i18n/locales/es.ts b/apps/desktop/src/i18n/locales/es.ts index 59cb3231ae..796f19de65 100644 --- a/apps/desktop/src/i18n/locales/es.ts +++ b/apps/desktop/src/i18n/locales/es.ts @@ -6546,6 +6546,9 @@ export default withEnglishFallback({ mcpCapabilityConnectionManagement: "Agregar/eliminar conexiones MCP", mcpCapabilityAllowed: "Permitido", mcpCapabilityBlocked: "Bloqueado", + mcpQueryTimeout: "Tiempo de espera de consulta (segundos)", + mcpQueryTimeoutDescription: "Déjalo vacío para heredar el tiempo de espera de la conexión. 0 = sin límite. Un valor positivo reemplaza la configuración de la conexión para cada consulta MCP.", + mcpQueryTimeoutInvalid: "Introduce un número entero no negativo de segundos o déjalo vacío para heredar la configuración de la conexión.", mcpCapabilityAlwaysEnforced: "Todos los modos siguen sujetos a la lista de conexiones permitidas, la protección de solo lectura, la protección de producción y los permisos de la cuenta de base de datos.", mcpPolicySaveFailed: "No se pudo guardar la política de acceso MCP: {error}", mcpPolicyLoadFailed: "No se pudo cargar la política de acceso MCP: {error}", diff --git a/apps/desktop/src/i18n/locales/it.ts b/apps/desktop/src/i18n/locales/it.ts index c036186605..3c7d7ecfd8 100644 --- a/apps/desktop/src/i18n/locales/it.ts +++ b/apps/desktop/src/i18n/locales/it.ts @@ -6546,6 +6546,9 @@ export default withEnglishFallback({ mcpCapabilityConnectionManagement: "Aggiunta/rimozione connessioni MCP", mcpCapabilityAllowed: "Consentito", mcpCapabilityBlocked: "Bloccato", + mcpQueryTimeout: "Timeout query (secondi)", + mcpQueryTimeoutDescription: "Lascia vuoto per ereditare il timeout della connessione. 0 = nessun limite. Un valore positivo sostituisce l'impostazione della connessione per ogni query MCP.", + mcpQueryTimeoutInvalid: "Inserisci un numero intero non negativo di secondi o lascia vuoto per ereditare l'impostazione della connessione.", mcpCapabilityAlwaysEnforced: "Tutte le modalità restano soggette all'elenco delle connessioni consentite, alla protezione di sola lettura, alla protezione di produzione e ai permessi dell'account database.", mcpPolicySaveFailed: "Impossibile salvare la policy di accesso MCP: {error}", mcpPolicyLoadFailed: "Impossibile caricare la policy di accesso MCP: {error}", diff --git a/apps/desktop/src/i18n/locales/ja.ts b/apps/desktop/src/i18n/locales/ja.ts index b8b441141b..ee7aa1efb6 100644 --- a/apps/desktop/src/i18n/locales/ja.ts +++ b/apps/desktop/src/i18n/locales/ja.ts @@ -6553,6 +6553,9 @@ export default withEnglishFallback({ mcpCapabilityConnectionManagement: "MCP接続の追加/削除", mcpCapabilityAllowed: "許可", mcpCapabilityBlocked: "禁止", + mcpQueryTimeout: "クエリタイムアウト（秒）", + mcpQueryTimeoutDescription: "空のままにすると接続のクエリタイムアウト設定を継承します。0 = 制限なし。正の値は接続設定を上書きし、すべての MCP クエリに適用されます。", + mcpQueryTimeoutInvalid: "秒数を負でない整数で入力するか、空にして接続設定を継承してください。", mcpCapabilityAlwaysEnforced: "すべてのモードで、接続許可リスト、接続の読み取り専用保護、本番環境保護、データベースアカウント権限が引き続き適用されます。", mcpPolicySaveFailed: "MCPアクセスポリシーを保存できませんでした: {error}", mcpPolicyLoadFailed: "MCPアクセスポリシーを読み込めませんでした: {error}", diff --git a/apps/desktop/src/i18n/locales/ko.ts b/apps/desktop/src/i18n/locales/ko.ts index 4b41242bce..50f8a5d575 100644 --- a/apps/desktop/src/i18n/locales/ko.ts +++ b/apps/desktop/src/i18n/locales/ko.ts @@ -6280,6 +6280,9 @@ export default withEnglishFallback({ mcpCapabilityConnectionManagement: "MCP 연결 추가/제거", mcpCapabilityAllowed: "허용됨", mcpCapabilityBlocked: "차단됨", + mcpQueryTimeout: "쿼리 시간 제한(초)", + mcpQueryTimeoutDescription: "비워 두면 연결의 쿼리 시간 제한 설정을 상속합니다. 0 = 제한 없음. 양수는 연결 설정을 재정의하여 모든 MCP 쿼리에 적용됩니다.", + mcpQueryTimeoutInvalid: "음수가 아닌 정수(초)를 입력하거나 비워 두어 연결 설정을 상속하십시오.", mcpCapabilityAlwaysEnforced: "모든 모드는 연결 허용 목록, 연결 읽기 전용 보호, 프로덕션 보호 및 데이터베이스 계정 권한의 적용을 받습니다.", mcpPolicySaveFailed: "MCP 접근 정책 저장 실패: {error}", mcpPolicyLoadFailed: "MCP 접근 정책 불러오기 실패: {error}", diff --git a/apps/desktop/src/i18n/locales/pt-BR.ts b/apps/desktop/src/i18n/locales/pt-BR.ts index b4330033bc..e47f87bc61 100644 --- a/apps/desktop/src/i18n/locales/pt-BR.ts +++ b/apps/desktop/src/i18n/locales/pt-BR.ts @@ -6548,6 +6548,9 @@ export default withEnglishFallback({ mcpCapabilityConnectionManagement: "Adicionar/remover conexões MCP", mcpCapabilityAllowed: "Permitido", mcpCapabilityBlocked: "Bloqueado", + mcpQueryTimeout: "Tempo limite de consulta (segundos)", + mcpQueryTimeoutDescription: "Deixe vazio para herdar o tempo limite da conexão. 0 = sem limite. Um valor positivo substitui a configuração da conexão para cada consulta MCP.", + mcpQueryTimeoutInvalid: "Insira um número inteiro não negativo de segundos ou deixe vazio para herdar a configuração da conexão.", mcpCapabilityAlwaysEnforced: "Todos os modos continuam sujeitos à lista de conexões permitidas, à proteção somente leitura, à proteção de produção e às permissões da conta do banco de dados.", mcpPolicySaveFailed: "Falha ao salvar a política de acesso MCP: {error}", mcpPolicyLoadFailed: "Falha ao carregar a política de acesso MCP: {error}", diff --git a/apps/desktop/src/i18n/locales/zh-CN.ts b/apps/desktop/src/i18n/locales/zh-CN.ts index b62049595d..c842d1bf33 100644 --- a/apps/desktop/src/i18n/locales/zh-CN.ts +++ b/apps/desktop/src/i18n/locales/zh-CN.ts @@ -6895,6 +6895,9 @@ export default withEnglishFallback({ mcpCapabilityConnectionManagement: "MCP 连接添加/删除", mcpCapabilityAllowed: "允许", mcpCapabilityBlocked: "禁止", + mcpQueryTimeout: "查询超时（秒）", + mcpQueryTimeoutDescription: "留空则继承连接的查询超时设置。0 = 不限时。正数会覆盖连接设置，对所有 MCP 查询生效。", + mcpQueryTimeoutInvalid: "请输入非负整数秒数，或留空以继承连接设置。", mcpCapabilityAlwaysEnforced: "所有模式仍受连接允许范围、连接只读、生产库保护和数据库账号权限约束。", mcpPolicySaveFailed: "保存 MCP 访问策略失败：{error}", mcpPolicyLoadFailed: "加载 MCP 访问策略失败：{error}", diff --git a/apps/desktop/src/i18n/locales/zh-TW.ts b/apps/desktop/src/i18n/locales/zh-TW.ts index 602bbe560b..22a2cfb200 100644 --- a/apps/desktop/src/i18n/locales/zh-TW.ts +++ b/apps/desktop/src/i18n/locales/zh-TW.ts @@ -5888,6 +5888,9 @@ export default withEnglishFallback({ mcpCapabilityConnectionManagement: "MCP 連線新增/刪除", mcpCapabilityAllowed: "允許", mcpCapabilityBlocked: "禁止", + mcpQueryTimeout: "查詢逾時（秒）", + mcpQueryTimeoutDescription: "留空則繼承連線的查詢逾時設定。0 = 不限時。正數會覆蓋連線設定，對所有 MCP 查詢生效。", + mcpQueryTimeoutInvalid: "請輸入非負整數秒數，或留空以繼承連線設定。", mcpCapabilityAlwaysEnforced: "所有模式仍受連線允許範圍、連線唯讀、生產庫保護和資料庫帳號權限約束。", mcpPolicySaveFailed: "儲存 MCP 存取策略失敗：{error}", mcpPolicyLoadFailed: "載入 MCP 存取策略失敗：{error}", diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index a0eb0dc398..6f0eda6392 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -246,6 +246,7 @@ export interface McpGlobalPolicy { allowDangerousSql: boolean; allowedConnectionIds: string[] | null; configured: boolean; + queryTimeoutSecs: number | null; } export interface SavedSqlSyncEntry { diff --git a/apps/desktop/src/stores/__tests__/settingsStore.spec.ts b/apps/desktop/src/stores/__tests__/settingsStore.spec.ts index c5c94e62e3..c705e97ae6 100644 --- a/apps/desktop/src/stores/__tests__/settingsStore.spec.ts +++ b/apps/desktop/src/stores/__tests__/settingsStore.spec.ts @@ -433,6 +433,7 @@ describe("normalizeMcpGlobalPolicy", () => { allowDangerousSql: false, allowedConnectionIds: null, configured: false, + queryTimeoutSecs: null, }); }); @@ -449,12 +450,20 @@ describe("normalizeMcpGlobalPolicy", () => { allowDangerousSql: true, allowedConnectionIds: ["connection-1", "connection-2"], configured: true, + queryTimeoutSecs: null, }); }); it("preserves an empty allowlist as deny all", () => { expect(normalizeMcpGlobalPolicy({ allowedConnectionIds: [] }).allowedConnectionIds).toEqual([]); }); + + it("round-trips queryTimeoutSecs null, undefined and positive numbers", () => { + expect(normalizeMcpGlobalPolicy({ queryTimeoutSecs: null }).queryTimeoutSecs).toBeNull(); + expect(normalizeMcpGlobalPolicy({ queryTimeoutSecs: undefined } as any).queryTimeoutSecs).toBeNull(); + expect(normalizeMcpGlobalPolicy({ queryTimeoutSecs: 0 }).queryTimeoutSecs).toBe(0); + expect(normalizeMcpGlobalPolicy({ queryTimeoutSecs: 300 }).queryTimeoutSecs).toBe(300); + }); }); describe("normalizeEditorSettings - continueOnErrorOnBatch", () => { @@ -667,6 +676,7 @@ describe("settingsStore MCP policy persistence", () => { allowDangerousSql: false, allowedConnectionIds: ["connection-1"], configured: true, + queryTimeoutSecs: null, }; store.mcpGlobalPolicy = previous; @@ -679,6 +689,7 @@ describe("settingsStore MCP policy persistence", () => { allowDangerousSql: false, allowedConnectionIds: [], configured: true, + queryTimeoutSecs: null, }); rejectSave(new Error("save failed")); diff --git a/apps/desktop/src/stores/settingsStore.ts b/apps/desktop/src/stores/settingsStore.ts index 71a2aeb2a8..72921c7552 100644 --- a/apps/desktop/src/stores/settingsStore.ts +++ b/apps/desktop/src/stores/settingsStore.ts @@ -51,6 +51,8 @@ export interface McpGlobalPolicy { allowDangerousSql: boolean; allowedConnectionIds: string[] | null; configured: boolean; + /** MCP query timeout override in seconds. null/undefined = inherit the connection; 0 = no limit. */ + queryTimeoutSecs: number | null; } export type DesktopIconTheme = "default" | "black"; @@ -90,15 +92,21 @@ export const DEFAULT_MCP_GLOBAL_POLICY: McpGlobalPolicy = { allowDangerousSql: false, allowedConnectionIds: null, configured: false, + queryTimeoutSecs: null, }; export function normalizeMcpGlobalPolicy(policy: Partial | null | undefined): McpGlobalPolicy { const allowedConnectionIds = policy?.allowedConnectionIds === null || policy?.allowedConnectionIds === undefined ? null : [...new Set(policy.allowedConnectionIds.filter((id): id is string => typeof id === "string" && id.trim().length > 0).map((id) => id.trim()))]; + // null / undefined / non-positive => null (inherit connection). 0 is preserved + // as an explicit "no limit" only here; the UI maps "no limit" <=> 0 and + // "inherit" <=> null. + const queryTimeoutSecs = policy?.queryTimeoutSecs === null || policy?.queryTimeoutSecs === undefined ? null : typeof policy.queryTimeoutSecs === "number" && Number.isFinite(policy.queryTimeoutSecs) && policy.queryTimeoutSecs >= 0 ? Math.round(policy.queryTimeoutSecs) : null; return { readOnly: policy?.readOnly === true, allowDangerousSql: policy?.allowDangerousSql === true, allowedConnectionIds, configured: policy?.configured === true, + queryTimeoutSecs, }; } @@ -1522,6 +1530,7 @@ export const useSettingsStore = defineStore("settings", () => { readOnly: next.readOnly, allowDangerousSql: next.allowDangerousSql, allowedConnectionIds: next.allowedConnectionIds, + queryTimeoutSecs: next.queryTimeoutSecs, }); } catch (error) { mcpGlobalPolicy.value = previous; diff --git a/crates/dbx-core/src/agent_tools.rs b/crates/dbx-core/src/agent_tools.rs index ba78fadae7..d254cef57d 100644 --- a/crates/dbx-core/src/agent_tools.rs +++ b/crates/dbx-core/src/agent_tools.rs @@ -2,11 +2,13 @@ use std::collections::HashMap; use std::sync::{Arc, Mutex as StdMutex, OnceLock, Weak}; use serde_json::json; +use serde_json::Value; use crate::agent_events::{ToolCall, ToolDefinition, ToolResult}; use crate::connection::AppState; use crate::db::vector_driver; use crate::models::connection::DatabaseType; +use crate::models::connection::{ConnectionConfig, SPANNER_MIN_QUERY_TIMEOUT_SECS}; use crate::query::QueryExecutionOptions; use crate::query_execution_sql::{build_explain_sql, supports_explain_plan, supports_sql_query, ExplainSqlOptions}; use crate::sql_dialect::{build_table_data_select_sql, TableDataSelectSqlOptions}; @@ -63,6 +65,33 @@ fn tool_uses_database(tool_name: &str) -> bool { ) } +/// Resolve the query timeout (seconds) for an agent/MCP query. +/// +/// Precedence: an explicit per-call `requested` value wins; otherwise inherit +/// the connection's `effective_query_timeout_secs()`; the legacy 30s survives +/// only when no connection config is available to resolve against. `Some(0)` +/// from either source means "no limit" (unlimited, and it bypasses the Spanner +/// floor — see `ConnectionConfig::effective_query_timeout_secs`). A finite +/// value on a Spanner connection gets the 120s floor. +pub fn agent_query_timeout_secs(requested: Option, connection: Option<&ConnectionConfig>) -> u64 { + match requested { + Some(0) => 0, + Some(n) => { + if connection.is_some_and(|config| config.db_type == DatabaseType::Spanner) { + n.max(SPANNER_MIN_QUERY_TIMEOUT_SECS) + } else { + n + } + } + None => connection.map_or(QUERY_TIMEOUT_SECS, ConnectionConfig::effective_query_timeout_secs), + } +} + +/// Legacy fallback query timeout for agent queries. Matches the current fixed +/// `Some(30)` behavior and remains the last resort when no connection config +/// can be resolved (that path fails the query at pool lookup anyway). +const QUERY_TIMEOUT_SECS: u64 = 30; + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct QueryCellWindow { offset: usize, @@ -833,10 +862,16 @@ async fn execute_execute_query( let client_session_id = tool_call.arguments.get("client_session_id").and_then(|v| v.as_str()).map(str::trim).filter(|v| !v.is_empty()); - // Execute query using existing infrastructure + // Execute query using existing infrastructure. Timeout resolves as + // per-call `timeout_secs` > connection effective timeout (see + // `agent_query_timeout_secs`); MCP injects the policy value as + // `timeout_secs` in the arguments. let options = QueryExecutionOptions { max_rows: Some(limit), - timeout_secs: Some(30), + timeout_secs: Some(agent_query_timeout_secs( + tool_call.arguments.get("timeout_secs").and_then(Value::as_u64), + connection_config.as_ref(), + )), client_session_id: client_session_id.map(str::to_string), ..Default::default() }; @@ -1085,8 +1120,18 @@ async fn execute_explain_query( } }; - // Execute the EXPLAIN query - let options = QueryExecutionOptions { max_rows: Some(100), timeout_secs: Some(30), ..Default::default() }; + // Execute the EXPLAIN query. Timeout resolves like the execute path + // (per-call `timeout_secs` > connection effective timeout) so the global + // MCP timeout override covers EXPLAIN too. + let connection_config = state.configs.read().await.get(connection_id).cloned(); + let options = QueryExecutionOptions { + max_rows: Some(100), + timeout_secs: Some(agent_query_timeout_secs( + tool_call.arguments.get("timeout_secs").and_then(Value::as_u64), + connection_config.as_ref(), + )), + ..Default::default() + }; let result = match crate::query::execute_sql_statement_with_options( state, connection_id, @@ -1280,6 +1325,50 @@ async fn resolve_chroma_collection_uuid( mod tests { use super::*; + #[test] + fn agent_query_timeout_resolves_by_precedence() { + // Explicit per-call value wins. + assert_eq!(agent_query_timeout_secs(Some(45), Some(&postgres_connection(60))), 45); + // 0 = unlimited bypasses the Spanner floor. + assert_eq!(agent_query_timeout_secs(Some(0), Some(&spanner_connection(60))), 0); + // Finite value on Spanner gets the 120s floor. + assert_eq!(agent_query_timeout_secs(Some(60), Some(&spanner_connection(60))), 120); + // No request + connection => inherit the connection's effective timeout. + assert_eq!(agent_query_timeout_secs(None, Some(&postgres_connection(60))), 60); + assert_eq!(agent_query_timeout_secs(None, Some(&postgres_connection(120))), 120); + assert_eq!(agent_query_timeout_secs(None, Some(&postgres_connection(0))), 0); + // No request, no connection => legacy 30s fallback. + assert_eq!(agent_query_timeout_secs(None, None), 30); + } + + fn postgres_connection(query_timeout_secs: u64) -> ConnectionConfig { + let mut config = test_connection(); + config.db_type = DatabaseType::Postgres; + config.query_timeout_secs = query_timeout_secs; + config + } + + fn spanner_connection(query_timeout_secs: u64) -> ConnectionConfig { + let mut config = test_connection(); + config.db_type = DatabaseType::Spanner; + config.query_timeout_secs = query_timeout_secs; + config + } + + fn test_connection() -> ConnectionConfig { + serde_json::from_value(json!({ + "id": "agent-timeout-test", + "name": "timeout-test", + "db_type": "postgres", + "host": "127.0.0.1", + "port": 5432, + "username": "", + "password": "", + "ssl": false, + })) + .expect("test connection config") + } + #[tokio::test] async fn database_tool_locks_are_connection_scoped_and_cancel_safe() { let first = connection_tool_lock("agent-tool-lock-a").lock_owned().await; diff --git a/crates/dbx-core/src/models/connection.rs b/crates/dbx-core/src/models/connection.rs index 1375b323f8..02114d6169 100644 --- a/crates/dbx-core/src/models/connection.rs +++ b/crates/dbx-core/src/models/connection.rs @@ -463,6 +463,15 @@ pub fn default_connect_timeout_secs() -> u64 { 10 } +/// Cloud Spanner schema changes are long-running operations rather than plain +/// statements. Measured against the real service, `CREATE INDEX` on an *empty* +/// table took 22.9-33.6s over seven runs, so the generic default fails +/// intermittently — and the failure is misleading, because the operation keeps +/// running server-side and completes, leaving a retry to report +/// `Duplicate name in schema`. Raising a floor mirrors what +/// `connection::agent_connect_timeout` already does for Access. +pub const SPANNER_MIN_QUERY_TIMEOUT_SECS: u64 = 120; + pub fn default_query_timeout_secs() -> u64 { 60 } @@ -824,21 +833,12 @@ impl ConnectionConfig { } } - /// Cloud Spanner schema changes are long-running operations rather than plain - /// statements. Measured against the real service, `CREATE INDEX` on an *empty* - /// table took 22.9-33.6s over seven runs, so the generic default fails - /// intermittently — and the failure is misleading, because the operation keeps - /// running server-side and completes, leaving a retry to report - /// `Duplicate name in schema`. Raising a floor mirrors what - /// `connection::agent_connect_timeout` already does for Access. - const SPANNER_MIN_QUERY_TIMEOUT_SECS: u64 = 120; - pub fn effective_query_timeout_secs(&self) -> u64 { if self.query_timeout_secs == 0 { // An explicit 0 is the UI's "no limit"; a floor must not impose one. return 0; } - let floor = if self.db_type == DatabaseType::Spanner { Self::SPANNER_MIN_QUERY_TIMEOUT_SECS } else { 1 }; + let floor = if self.db_type == DatabaseType::Spanner { SPANNER_MIN_QUERY_TIMEOUT_SECS } else { 1 }; self.query_timeout_secs.max(floor) } diff --git a/crates/dbx-core/src/storage.rs b/crates/dbx-core/src/storage.rs index c33d0c9a75..a908741685 100644 --- a/crates/dbx-core/src/storage.rs +++ b/crates/dbx-core/src/storage.rs @@ -224,6 +224,8 @@ pub struct McpGlobalPolicy { #[serde(default)] pub allow_dangerous_sql: bool, pub allowed_connection_ids: Option>, + #[serde(default)] + pub query_timeout_secs: Option, } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] @@ -233,6 +235,8 @@ pub struct McpGlobalPolicyState { pub read_only: bool, pub allow_dangerous_sql: bool, pub allowed_connection_ids: Option>, + #[serde(default)] + pub query_timeout_secs: Option, } impl McpGlobalPolicyState { @@ -241,6 +245,7 @@ impl McpGlobalPolicyState { read_only: self.read_only, allow_dangerous_sql: self.allow_dangerous_sql, allowed_connection_ids: self.allowed_connection_ids.clone(), + query_timeout_secs: self.query_timeout_secs, } } } @@ -1639,6 +1644,7 @@ impl Storage { read_only: policy.read_only, allow_dangerous_sql: policy.allow_dangerous_sql, allowed_connection_ids: policy.allowed_connection_ids, + query_timeout_secs: policy.query_timeout_secs, }); }; let settings = serde_json::from_str::>(&json) @@ -1650,6 +1656,7 @@ impl Storage { read_only: policy.read_only, allow_dangerous_sql: policy.allow_dangerous_sql, allowed_connection_ids: policy.allowed_connection_ids, + query_timeout_secs: policy.query_timeout_secs, }); }; let policy = serde_json::from_value::(value.clone()) @@ -1659,6 +1666,7 @@ impl Storage { read_only: policy.read_only, allow_dangerous_sql: policy.allow_dangerous_sql, allowed_connection_ids: policy.allowed_connection_ids, + query_timeout_secs: policy.query_timeout_secs, }) }) .await; @@ -6020,6 +6028,7 @@ mod tests { read_only: false, allow_dangerous_sql: false, allowed_connection_ids: None, + query_timeout_secs: None, } ); @@ -6029,6 +6038,7 @@ mod tests { read_only: true, allow_dangerous_sql: true, allowed_connection_ids: Some(vec!["conn-1".to_string(), "conn-2".to_string()]), + query_timeout_secs: Some(120), }) .await .unwrap(); @@ -6040,6 +6050,7 @@ mod tests { read_only: true, allow_dangerous_sql: true, allowed_connection_ids: Some(vec!["conn-1".to_string(), "conn-2".to_string()]), + query_timeout_secs: Some(120), } ); assert_eq!(storage.load_password_hash().await.unwrap().as_deref(), Some("preserved")); @@ -6047,6 +6058,7 @@ mod tests { assert_eq!(settings[MCP_GLOBAL_POLICY_KEY]["readOnly"], true); assert_eq!(settings[MCP_GLOBAL_POLICY_KEY]["allowDangerousSql"], true); assert_eq!(settings[MCP_GLOBAL_POLICY_KEY]["allowedConnectionIds"][0], "conn-1"); + assert_eq!(settings[MCP_GLOBAL_POLICY_KEY]["queryTimeoutSecs"], 120); assert!(settings[MCP_GLOBAL_POLICY_KEY].get("configured").is_none()); storage.save_desktop_settings(&DesktopSettings::default()).await.unwrap(); @@ -6119,6 +6131,7 @@ mod tests { let policy = storage.load_mcp_global_policy().await.unwrap(); assert!(policy.configured); assert!(!policy.allow_dangerous_sql); + assert_eq!(policy.query_timeout_secs, None); } #[tokio::test] @@ -6134,6 +6147,7 @@ mod tests { read_only: false, allow_dangerous_sql: false, allowed_connection_ids: Some(vec![kept.id.clone()]), + query_timeout_secs: None, }) .await .unwrap(); @@ -6157,6 +6171,7 @@ mod tests { read_only: true, allow_dangerous_sql: false, allowed_connection_ids: None, + query_timeout_secs: None, }) .await .unwrap(); diff --git a/crates/dbx-mcp/src/backend.rs b/crates/dbx-mcp/src/backend.rs index 9c3c9b8c3a..7f387bad03 100644 --- a/crates/dbx-mcp/src/backend.rs +++ b/crates/dbx-mcp/src/backend.rs @@ -822,6 +822,15 @@ impl DbxBackend for WebBackend { if let Some(client_session_id) = arguments.get("client_session_id").and_then(Value::as_str) { body["clientSessionId"] = json!(client_session_id); } + // The query timeout is resolved once here and carried on every + // request: policy > connection effective timeout (see + // `agent_query_timeout_secs`). `arguments["timeout_secs"]` holds the + // global MCP policy value injected by the server; an absent value + // means inherit the connection's setting. + body["timeoutSecs"] = json!(agent_tools::agent_query_timeout_secs( + arguments.get("timeout_secs").and_then(Value::as_u64), + Some(connection), + )); let response = self.request(reqwest::Method::POST, "/api/query/execute", Some(body)).await?; let query_result: dbx_core::db::QueryResult = response.json().await.map_err(|error| format!("Invalid query response: {error}"))?; @@ -870,7 +879,7 @@ impl DbxBackend for WebBackend { database: &str, sql: &str, _max_rows: Option, - _timeout_secs: Option, + timeout_secs: Option, ) -> Result { if connection.db_type == DatabaseType::MongoDb { return Err("MongoDB shell commands in DBX Web mode are not implemented by the Rust CLI yet.".to_string()); @@ -879,7 +888,7 @@ impl DbxBackend for WebBackend { self.request( reqwest::Method::POST, "/api/query/execute", - Some(json!({ "connectionId": connection.id, "database": database, "sql": sql })), + Some(json!({ "connectionId": connection.id, "database": database, "sql": sql, "timeoutSecs": agent_tools::agent_query_timeout_secs(timeout_secs, Some(connection)) })), ) .await? .json() @@ -1855,7 +1864,13 @@ mod tests { }; fn policy_state(configured: bool, read_only: bool) -> McpGlobalPolicyState { - McpGlobalPolicyState { configured, read_only, allow_dangerous_sql: false, allowed_connection_ids: None } + McpGlobalPolicyState { + configured, + read_only, + allow_dangerous_sql: false, + allowed_connection_ids: None, + query_timeout_secs: None, + } } #[test] @@ -1961,6 +1976,153 @@ mod tests { assert_eq!(paths.get("web-db"), Some(&vec!["Project".to_string()])); } + #[test] + fn agent_query_timeout_web_policy_applies_over_connection_effective_timeout() { + // Unset policy inherits the connection's effective timeout + let mut conn = new_connection_config( + "timeout".to_string(), + "timeout".to_string(), + DatabaseType::Postgres, + "localhost".to_string(), + 5432, + String::new(), + String::new(), + None, + false, + None, + ) + .unwrap(); + conn.query_timeout_secs = 60; + assert_eq!(agent_tools::agent_query_timeout_secs(None, Some(&conn)), 60); + + conn.query_timeout_secs = 0; + assert_eq!(agent_tools::agent_query_timeout_secs(None, Some(&conn)), 0); + + // Policy 300 overrides both a 60 and a 0 connection. + assert_eq!(agent_tools::agent_query_timeout_secs(Some(300), Some(&conn)), 300); + + // Spanner floor applies to a finite policy but not to 0. + let mut spanner = new_connection_config( + "spanner".to_string(), + "spanner".to_string(), + DatabaseType::Spanner, + "localhost".to_string(), + 5432, + String::new(), + String::new(), + None, + false, + None, + ) + .unwrap(); + spanner.query_timeout_secs = 60; + assert_eq!(agent_tools::agent_query_timeout_secs(Some(60), Some(&spanner)), 120); + assert_eq!(agent_tools::agent_query_timeout_secs(Some(0), Some(&spanner)), 0); + assert_eq!(agent_tools::agent_query_timeout_secs(None, Some(&spanner)), 120); + } + + #[tokio::test] + async fn web_execute_agent_tool_carries_resolved_timeout_secs_in_body() { + 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 || { + for _ in 0..2 { + 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]; + let header_end = loop { + let count = stream.read(&mut buffer).unwrap(); + request.extend_from_slice(&buffer[..count]); + if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") { + break position + 4; + } + }; + let headers = String::from_utf8_lossy(&request[..header_end]); + let content_length = headers + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length").then_some(value.trim()) + }) + .unwrap() + .parse::() + .unwrap(); + while request.len() < header_end + content_length { + let count = stream.read(&mut buffer).unwrap(); + request.extend_from_slice(&buffer[..count]); + } + let request = String::from_utf8(request).unwrap(); + let body = request[header_end..header_end + content_length].to_string(); + request_sender.send((request.lines().next().unwrap().to_string(), body)).unwrap(); + + let response_body = r#"{"columns":["?"],"column_types":[],"column_sortables":[],"rows":[["1"]],"affected_rows":0,"execution_time_ms":0,"truncated":false,"has_more":false}"#; + write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", + response_body.len(), + response_body + ) + .unwrap(); + } + }); + + let backend = WebBackend::new(format!("http://{address}"), String::new()).unwrap(); + backend.auth.lock().await.checked = true; + let mut connection = new_connection_config( + "web-timeout".to_string(), + "web-timeout".to_string(), + DatabaseType::Postgres, + "localhost".to_string(), + 5432, + String::new(), + String::new(), + None, + false, + None, + ) + .unwrap(); + connection.query_timeout_secs = 60; + // ensure_connected only POSTs /api/connection/connect when the backend has + // no cached config; pre-seed the cache so the mock needs to answer only the + // /api/query/execute request. + backend.connected.lock().await.insert(connection.id.clone(), connection.clone()); + + // No policy argument -> inherit the connection's effective timeout (60). + let result = backend + .execute_agent_tool( + &connection, + "postgres", + "execute_query", + json!({ "sql": "SELECT 1", "limit": 10 }), + AgentSqlPermissions { allow_writes: false, allow_dangerous: false, confirmed_write_sql: None }, + ) + .await; + assert!(!result.is_error, "{:?}", result); + let (request_line, body) = request_receiver.recv().unwrap(); + assert_eq!(request_line, "POST /api/query/execute HTTP/1.1"); + let request: Value = serde_json::from_str(&body).unwrap(); + assert_eq!(request["timeoutSecs"], 60); + + // Policy argument 300 overrides the connection. + let result = backend + .execute_agent_tool( + &connection, + "postgres", + "execute_query", + json!({ "sql": "SELECT 1", "limit": 10, "timeout_secs": 300 }), + AgentSqlPermissions { allow_writes: false, allow_dangerous: false, confirmed_write_sql: None }, + ) + .await; + assert!(!result.is_error, "{:?}", result); + + server.join().unwrap(); + let (_request_line, second_body) = request_receiver.recv().unwrap(); + let second_request: Value = serde_json::from_str(&second_body).unwrap(); + assert_eq!(second_request["timeoutSecs"], 300); + } + #[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 0ca7da46ac..2c3a787bee 100644 --- a/crates/dbx-mcp/src/server.rs +++ b/crates/dbx-mcp/src/server.rs @@ -494,6 +494,14 @@ impl DbxMcpServer { if let Some(limit) = request.cell_char_limit { arguments["cell_char_limit"] = json!(limit); } + // Surface the global MCP query timeout as a per-query argument. It is + // re-read on every tool call, so a settings-page change takes effect on + // the next MCP query with no client-config regeneration. A future + // per-call `timeout_secs` tool parameter can slot in above this by not + // being clobbered when set by the model. + if let Some(secs) = resolved.policy.query_timeout_secs { + arguments["timeout_secs"] = json!(secs); + } let result = self.backend.execute_agent_tool(connection, &database, "execute_query", arguments, permissions).await; agent_result(result) @@ -1905,7 +1913,12 @@ mod tests { fn local_mongo_aggregate_cannot_write_to_a_production_database() { let mut mongo = connection("mongo", "mongo", "mongodb", "staging"); mongo.production_databases = vec!["production".to_string()]; - let policy = McpGlobalPolicy { read_only: false, allow_dangerous_sql: true, allowed_connection_ids: None }; + let policy = McpGlobalPolicy { + read_only: false, + allow_dangerous_sql: true, + allowed_connection_ids: None, + query_timeout_secs: None, + }; let error = validate_mongo_command( &mongo, @@ -1927,9 +1940,24 @@ mod tests { let source = r#"db.runCommand({ping: 1})"#; for policy in [ - McpGlobalPolicy { read_only: true, allow_dangerous_sql: false, allowed_connection_ids: None }, - McpGlobalPolicy { read_only: false, allow_dangerous_sql: false, allowed_connection_ids: None }, - McpGlobalPolicy { read_only: false, allow_dangerous_sql: true, allowed_connection_ids: None }, + McpGlobalPolicy { + read_only: true, + allow_dangerous_sql: false, + allowed_connection_ids: None, + query_timeout_secs: None, + }, + McpGlobalPolicy { + read_only: false, + allow_dangerous_sql: false, + allowed_connection_ids: None, + query_timeout_secs: None, + }, + McpGlobalPolicy { + read_only: false, + allow_dangerous_sql: true, + allowed_connection_ids: None, + query_timeout_secs: None, + }, ] { let error = validate_mongo_command(&mongo, &policy, "staging", source).unwrap_err(); assert!(result_text(&error).contains("SQL_BLOCKED")); @@ -1958,8 +1986,12 @@ mod tests { assert_eq!(normalize_confirmed_write_sql(Some(" \n ".to_string())), None); let read_only = ConnectionConfig { read_only: true, ..connection("readonly", "readonly", "postgres", "app") }; - let writable_policy = - McpGlobalPolicy { read_only: false, allow_dangerous_sql: true, allowed_connection_ids: None }; + let writable_policy = McpGlobalPolicy { + read_only: false, + allow_dangerous_sql: true, + allowed_connection_ids: None, + query_timeout_secs: None, + }; let read_only_error = validate_sql_policy(&read_only, &writable_policy, "app", "DELETE FROM sessions", false).unwrap_err(); assert!(result_text(&read_only_error).contains("CONNECTION_READ_ONLY")); @@ -2009,7 +2041,12 @@ mod tests { // An exact confirmation remains available as a narrowing constraint, // but cannot turn the central safe-write policy into full access. let _guard = EnvGuard::set("DBX_MCP_CONFIRMED_WRITE_SQL", "CREATE TABLE metrics (id INT)"); - let safe_write = McpGlobalPolicy { read_only: false, allow_dangerous_sql: false, allowed_connection_ids: None }; + let safe_write = McpGlobalPolicy { + read_only: false, + allow_dangerous_sql: false, + allowed_connection_ids: None, + query_timeout_secs: None, + }; let permissions = mcp_permissions(&connection, &safe_write); assert!(!permissions.allow_dangerous, "confirmed SQL must NOT elevate allow_dangerous for Redis/Mongo paths"); assert!(permissions.allow_writes); @@ -2031,7 +2068,12 @@ mod tests { // Full access still permits an exact confirmed DDL statement and keeps // the binding for the execution-time anti-replay check. let _guard = EnvGuard::set("DBX_MCP_CONFIRMED_WRITE_SQL", "CREATE TABLE metrics (id INT)"); - let full_access = McpGlobalPolicy { read_only: false, allow_dangerous_sql: true, allowed_connection_ids: None }; + let full_access = McpGlobalPolicy { + read_only: false, + allow_dangerous_sql: true, + allowed_connection_ids: None, + query_timeout_secs: None, + }; let permissions = validate_sql_policy(&connection, &full_access, "app", "CREATE TABLE metrics (id INT)", false).unwrap(); assert!(permissions.allow_writes); @@ -2040,8 +2082,12 @@ mod tests { // A global read-only policy is still authoritative even for the exact // confirmed statement, while read queries remain available. - let read_only_policy = - McpGlobalPolicy { read_only: true, allow_dangerous_sql: false, allowed_connection_ids: None }; + let read_only_policy = McpGlobalPolicy { + read_only: true, + allow_dangerous_sql: false, + allowed_connection_ids: None, + query_timeout_secs: None, + }; let error = validate_sql_policy(&connection, &read_only_policy, "app", "CREATE TABLE metrics (id INT)", false) .unwrap_err(); assert!(result_text(&error).contains("MCP_READ_ONLY"), "confirmed SQL must not bypass global read_only"); @@ -2061,7 +2107,12 @@ mod tests { #[test] fn use_statements_require_a_session() { let starrocks = connection("sr", "sr", "starrocks", "default_catalog"); - let policy = McpGlobalPolicy { read_only: false, allow_dangerous_sql: false, allowed_connection_ids: None }; + let policy = McpGlobalPolicy { + read_only: false, + allow_dangerous_sql: false, + allowed_connection_ids: None, + query_timeout_secs: None, + }; let blocked = validate_sql_policy(&starrocks, &policy, "default_catalog", "USE analytics", false).unwrap_err(); assert!(result_text(&blocked).contains("SQL_BLOCKED")); diff --git a/crates/dbx-mcp/tests/local.rs b/crates/dbx-mcp/tests/local.rs index 45676bd438..3e3dc73ac7 100644 --- a/crates/dbx-mcp/tests/local.rs +++ b/crates/dbx-mcp/tests/local.rs @@ -99,6 +99,7 @@ async fn duplicate_connection_preserves_secrets_ssh_and_sidebar_group() { read_only: false, allow_dangerous_sql: false, allowed_connection_ids: None, + query_timeout_secs: None, }) .await .expect("enable MCP connection management"); @@ -233,6 +234,7 @@ async fn duplicate_connection_preserves_secrets_ssh_and_sidebar_group() { read_only: true, allow_dangerous_sql: false, allowed_connection_ids: None, + query_timeout_secs: None, }) .await .expect("enable MCP read-only policy"); diff --git a/crates/dbx-mcp/tests/protocol.rs b/crates/dbx-mcp/tests/protocol.rs index ed5f8c5c30..dcaeae180c 100644 --- a/crates/dbx-mcp/tests/protocol.rs +++ b/crates/dbx-mcp/tests/protocol.rs @@ -1,4 +1,4 @@ -use std::{collections::HashMap, sync::Arc}; +use std::{collections::HashMap, sync::Arc, sync::Mutex}; use async_trait::async_trait; use dbx_core::{ @@ -111,6 +111,123 @@ impl DbxBackend for PolicyBackend { } } +struct CapturingBackend { + policy: McpGlobalPolicy, + connections: Vec, + /// Records the arguments passed to `execute_agent_tool` so tests can assert + /// what the server injected (e.g. `timeout_secs` from the MCP policy). + calls: Mutex>, +} + +#[async_trait] +impl DbxBackend for CapturingBackend { + async fn load_mcp_global_policy(&self) -> Result { + Ok(self.policy.clone()) + } + + async fn load_connections(&self) -> Result, String> { + Ok(self.connections.clone()) + } + + async fn execute_agent_tool( + &self, + _connection: &ConnectionConfig, + _database: &str, + tool_name: &str, + arguments: Value, + _permissions: AgentSqlPermissions, + ) -> ToolResult { + self.calls.lock().expect("lock captures").push(arguments); + ToolResult { + tool_call_id: "capture-test".to_string(), + tool_name: tool_name.to_string(), + content: "ok".to_string(), + is_error: false, + explain_data: None, + } + } + + async fn add_connection_for_mcp(&self, config: ConnectionConfig) -> Result { + Ok(config) + } + + async fn duplicate_connection_for_mcp( + &self, + _source_id: &str, + _copy_id: &str, + _copy_name: &str, + ) -> Result { + Err("not exercised".to_string()) + } + + async fn remove_connection_for_mcp(&self, _connection_id: &str) -> Result { + Ok(true) + } +} + +/// Drives one `dbx_execute_query` call against a `CapturingBackend`, returning +/// the captured arguments for assertion. +async fn captured_query_arguments(backend: Arc, sql_arguments: Value) -> Vec { + let (server_transport, client_transport) = tokio::io::duplex(16 * 1024); + let server = DbxMcpServer::with_runtime_options(backend.clone(), 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 _ = client + .peer() + .call_tool( + CallToolRequestParams::new("dbx_execute_query") + .with_arguments(sql_arguments.as_object().cloned().unwrap_or_else(Map::new)), + ) + .await + .expect("call execute_query"); + client.cancel().await.expect("close MCP client"); + server_task.abort(); + backend.calls.lock().expect("lock captures").clone() +} + +/// Regression for the MCP global query-timeout override: the server must inject +/// `timeout_secs` into the `dbx_execute_query` arguments whenever the persisted +/// policy carries a value, and must NOT inject it when the policy inherits the +/// connection (None). This is the native boundary that turns the settings-page +/// value into a per-query argument (server.rs `if let Some(secs) = ...`). +#[tokio::test] +async fn execute_query_injects_and_omits_timeout_secs_from_policy() { + // Case 1: a positive policy timeout is injected on every call. + let backend = Arc::new(CapturingBackend { + policy: McpGlobalPolicy { + read_only: false, + allow_dangerous_sql: false, + allowed_connection_ids: None, + query_timeout_secs: Some(300), + }, + connections: vec![test_connection("scoped", "shared-db")], + calls: Mutex::new(Vec::new()), + }); + let captured = captured_query_arguments(backend, json!({ "connection_id": "scoped", "sql": "SELECT 1" })).await; + assert_eq!(captured.len(), 1, "expected one captured execute_query call"); + assert_eq!(captured[0]["timeout_secs"], json!(300), "policy timeout must be injected"); + + // Case 2: an inheriting policy (None) omits the key so the connection-level + // default is honoured rather than overridden by a stale value. + let backend = Arc::new(CapturingBackend { + policy: McpGlobalPolicy { + read_only: false, + allow_dangerous_sql: false, + allowed_connection_ids: None, + query_timeout_secs: None, + }, + connections: vec![test_connection("scoped", "shared-db")], + calls: Mutex::new(Vec::new()), + }); + let captured = captured_query_arguments(backend, json!({ "connection_id": "scoped", "sql": "SELECT 1" })).await; + assert_eq!(captured.len(), 1, "expected one captured execute_query call"); + assert!( + !captured[0].get("timeout_secs").is_some(), + "no timeout_secs must be injected when the policy inherits the connection: {}", + captured[0] + ); +} + fn test_connection(id: &str, name: &str) -> ConnectionConfig { serde_json::from_value(json!({ "id": id, @@ -193,6 +310,7 @@ async fn enforces_global_connection_scope_and_read_only_policy() { read_only: true, allow_dangerous_sql: false, allowed_connection_ids: Some(vec!["allowed".to_string(), "allowed-staging".to_string()]), + query_timeout_secs: None, }, connections: vec![ test_connection("allowed", "shared-db"), @@ -261,7 +379,12 @@ async fn read_only_policy_blocks_write_capable_sql_the_keyword_scan_misses() { (true, "SELECT * FROM users FOR SHARE"), ] { let backend = PolicyBackend { - policy: McpGlobalPolicy { read_only: true, allow_dangerous_sql, allowed_connection_ids: None }, + policy: McpGlobalPolicy { + read_only: true, + allow_dangerous_sql, + allowed_connection_ids: None, + query_timeout_secs: None, + }, connections: vec![mysql_connection("mysql", "reporting")], group_paths: Ok(HashMap::new()), }; @@ -301,7 +424,12 @@ async fn read_only_policy_allows_read_only_show_statements() { ] { let connection_id = connection.id.clone(); let backend = PolicyBackend { - policy: McpGlobalPolicy { read_only: true, allow_dangerous_sql: false, allowed_connection_ids: None }, + policy: McpGlobalPolicy { + read_only: true, + allow_dangerous_sql: false, + allowed_connection_ids: None, + query_timeout_secs: None, + }, connections: vec![connection], group_paths: Ok(HashMap::new()), }; diff --git a/crates/dbx-web/src/routes/connection.rs b/crates/dbx-web/src/routes/connection.rs index 594b486ab8..ba8ca432d2 100644 --- a/crates/dbx-web/src/routes/connection.rs +++ b/crates/dbx-web/src/routes/connection.rs @@ -1227,6 +1227,7 @@ mod tests { read_only: false, allow_dangerous_sql: false, allowed_connection_ids: Some(vec![existing.id.clone()]), + query_timeout_secs: None, }) .await .unwrap(); @@ -1299,6 +1300,7 @@ mod tests { read_only: false, allow_dangerous_sql: false, allowed_connection_ids: Some(vec![removed.id.clone()]), + query_timeout_secs: None, }) .await .unwrap(); @@ -1329,6 +1331,7 @@ mod tests { read_only: true, allow_dangerous_sql: false, allowed_connection_ids: None, + query_timeout_secs: None, }) .await .unwrap(); diff --git a/crates/dbx-web/src/routes/mongo.rs b/crates/dbx-web/src/routes/mongo.rs index 93684642c0..e1110ac87e 100644 --- a/crates/dbx-web/src/routes/mongo.rs +++ b/crates/dbx-web/src/routes/mongo.rs @@ -986,6 +986,7 @@ mod tests { read_only: false, allow_dangerous_sql: false, allowed_connection_ids: Some(vec![connection.id.clone()]), + query_timeout_secs: None, }; state.app.storage.save_mcp_global_policy(&writable_policy).await.unwrap(); @@ -1032,6 +1033,7 @@ mod tests { read_only: false, allow_dangerous_sql: true, allowed_connection_ids: Some(vec!["different-connection".to_string()]), + query_timeout_secs: None, }) .await .unwrap(); @@ -1056,6 +1058,7 @@ mod tests { read_only: false, allow_dangerous_sql: false, allowed_connection_ids: Some(vec![connection.id.clone()]), + query_timeout_secs: None, }) .await .unwrap(); @@ -1085,6 +1088,7 @@ mod tests { read_only: false, allow_dangerous_sql: false, allowed_connection_ids: Some(vec![connection.id.clone()]), + query_timeout_secs: None, }) .await .unwrap(); diff --git a/crates/dbx-web/src/routes/mq.rs b/crates/dbx-web/src/routes/mq.rs index 60f975e27f..35ad676778 100644 --- a/crates/dbx-web/src/routes/mq.rs +++ b/crates/dbx-web/src/routes/mq.rs @@ -1586,6 +1586,7 @@ mod tests { read_only: false, allow_dangerous_sql: true, allowed_connection_ids: Some(vec![connection_id.to_string()]), + query_timeout_secs: None, } } diff --git a/src-tauri/src/commands/mcp_bridge.rs b/src-tauri/src/commands/mcp_bridge.rs index f8556376c4..ebd5083c46 100644 --- a/src-tauri/src/commands/mcp_bridge.rs +++ b/src-tauri/src/commands/mcp_bridge.rs @@ -420,19 +420,29 @@ mod tests { #[test] fn mcp_allowlist_distinguishes_all_subset_and_none() { - let all = McpGlobalPolicy { read_only: false, allow_dangerous_sql: false, allowed_connection_ids: None }; + let all = McpGlobalPolicy { + read_only: false, + allow_dangerous_sql: false, + allowed_connection_ids: None, + query_timeout_secs: None, + }; assert!(ensure_connection_in_mcp_scope(&all, "conn-1").is_ok()); let subset = McpGlobalPolicy { read_only: false, allow_dangerous_sql: false, allowed_connection_ids: Some(vec!["conn-1".to_string()]), + query_timeout_secs: None, }; assert!(ensure_connection_in_mcp_scope(&subset, "conn-1").is_ok()); assert!(ensure_connection_in_mcp_scope(&subset, "conn-2").unwrap_err().starts_with("CONNECTION_OUT_OF_SCOPE:")); - let none = - McpGlobalPolicy { read_only: false, allow_dangerous_sql: false, allowed_connection_ids: Some(Vec::new()) }; + let none = McpGlobalPolicy { + read_only: false, + allow_dangerous_sql: false, + allowed_connection_ids: Some(Vec::new()), + query_timeout_secs: None, + }; assert!(ensure_connection_in_mcp_scope(&none, "conn-1").is_err()); }