From 6ef7eba566639ed271d08f6580db9a9430227a02 Mon Sep 17 00:00:00 2001
From: t8y2 <1156263951@qq.com>
Date: Wed, 3 Jun 2026 02:39:07 +0800
Subject: [PATCH] fix(mongodb): allow legacy fallback after native timeout

---
 apps/desktop/src/lib/connectionAttemptTimeout.ts    |  7 +++++--
 crates/dbx-core/src/db/mongo_driver.rs              | 13 +++++++++----
 packages/app-tests/connectionAttemptTimeout.test.ts |  8 ++++++++
 3 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/apps/desktop/src/lib/connectionAttemptTimeout.ts b/apps/desktop/src/lib/connectionAttemptTimeout.ts
index 26a74e272..b2fc8de89 100644
--- a/apps/desktop/src/lib/connectionAttemptTimeout.ts
+++ b/apps/desktop/src/lib/connectionAttemptTimeout.ts
@@ -1,6 +1,7 @@
 import type { ConnectionConfig } from "@/types/database";
 
 export const CONNECTION_ATTEMPT_TIMEOUT_BUFFER_MS = 2_000;
+export const MONGO_LEGACY_FALLBACK_TIMEOUT_BUFFER_MS = 30_000;
 const DEFAULT_CONNECT_TIMEOUT_SECS = 5;
 
 function positiveSeconds(value: unknown, fallback: number): number {
@@ -8,7 +9,8 @@ function positiveSeconds(value: unknown, fallback: number): number {
 }
 
 export function connectionAttemptTimeoutMs(
-  config: Pick<ConnectionConfig, "connect_timeout_secs" | "ssh_enabled" | "ssh_connect_timeout_secs" | "ssh_tunnels">,
+  config: Pick<ConnectionConfig, "connect_timeout_secs" | "ssh_enabled" | "ssh_connect_timeout_secs" | "ssh_tunnels"> &
+    Partial<Pick<ConnectionConfig, "db_type">>,
 ): number {
   const timeouts = [positiveSeconds(config.connect_timeout_secs, DEFAULT_CONNECT_TIMEOUT_SECS)];
   if (config.ssh_enabled) {
@@ -17,7 +19,8 @@ export function connectionAttemptTimeoutMs(
       timeouts.push(positiveSeconds(tunnel.connect_timeout_secs, DEFAULT_CONNECT_TIMEOUT_SECS));
     }
   }
-  return Math.ceil(Math.max(...timeouts) * 1000 + CONNECTION_ATTEMPT_TIMEOUT_BUFFER_MS);
+  const fallbackBuffer = config.db_type === "mongodb" ? MONGO_LEGACY_FALLBACK_TIMEOUT_BUFFER_MS : 0;
+  return Math.ceil(Math.max(...timeouts) * 1000 + CONNECTION_ATTEMPT_TIMEOUT_BUFFER_MS + fallbackBuffer);
 }
 
 export function connectionAttemptTimeoutMessage(timeoutMs: number): string {
diff --git a/crates/dbx-core/src/db/mongo_driver.rs b/crates/dbx-core/src/db/mongo_driver.rs
index 6aa7c0cc1..ffb476648 100644
--- a/crates/dbx-core/src/db/mongo_driver.rs
+++ b/crates/dbx-core/src/db/mongo_driver.rs
@@ -1,5 +1,6 @@
 use mongodb::{
     bson::{doc, oid::ObjectId, Bson, Document},
+    options::ClientOptions,
     Client,
 };
 use serde::{Deserialize, Serialize};
@@ -15,16 +16,20 @@ pub struct MongoDocumentResult {
 
 pub async fn connect(url: &str, timeout: Duration) -> Result<Client, String> {
     with_connection_timeout("MongoDB", timeout, async {
-        Client::with_uri_str(url).await.map_err(|e| format!("MongoDB connection failed: {e}"))
+        let mut options = ClientOptions::parse(url).await.map_err(|e| format!("MongoDB connection failed: {e}"))?;
+        options.connect_timeout = Some(timeout);
+        options.server_selection_timeout = Some(timeout);
+        Client::with_options(options).map_err(|e| format!("MongoDB connection failed: {e}"))
     })
     .await
 }
 
-pub async fn test_connection(client: &Client, timeout: Duration, database: Option<&str>) -> Result<(), String> {
+pub async fn test_connection(client: &Client, _timeout: Duration, database: Option<&str>) -> Result<(), String> {
     let database = database.map(str::trim).filter(|value| !value.is_empty()).unwrap_or("admin");
-    tokio::time::timeout(timeout, client.database(database).run_command(doc! { "ping": 1 }))
+    client
+        .database(database)
+        .run_command(doc! { "ping": 1 })
         .await
-        .map_err(|_| format!("MongoDB connection timed out ({}s)", timeout.as_secs()))?
         .map(|_| ())
         .map_err(|e| format!("MongoDB connection failed: {e}"))
 }
diff --git a/packages/app-tests/connectionAttemptTimeout.test.ts b/packages/app-tests/connectionAttemptTimeout.test.ts
index f9b09419e..663d8b0a1 100644
--- a/packages/app-tests/connectionAttemptTimeout.test.ts
+++ b/packages/app-tests/connectionAttemptTimeout.test.ts
@@ -2,6 +2,7 @@ import { strict as assert } from "node:assert";
 import test from "node:test";
 import {
   CONNECTION_ATTEMPT_TIMEOUT_BUFFER_MS,
+  MONGO_LEGACY_FALLBACK_TIMEOUT_BUFFER_MS,
   connectionAttemptTimeoutMessage,
   connectionAttemptTimeoutMs,
 } from "../../apps/desktop/src/lib/connectionAttemptTimeout.ts";
@@ -27,6 +28,13 @@ test("honors slower SSH tunnel connection timeouts", () => {
   );
 });
 
+test("allows MongoDB legacy agent fallback after native driver timeout", () => {
+  assert.equal(
+    connectionAttemptTimeoutMs({ db_type: "mongodb", connect_timeout_secs: 5 }),
+    5_000 + CONNECTION_ATTEMPT_TIMEOUT_BUFFER_MS + MONGO_LEGACY_FALLBACK_TIMEOUT_BUFFER_MS,
+  );
+});
+
 test("formats connection attempt timeout messages", () => {
   assert.match(connectionAttemptTimeoutMessage(7_001), /timed out after 8s/);
   assert.match(connectionAttemptTimeoutMessage(7_001), /VPN/);