From 35535be419e74b7859727323600ad0c9dbdbb432 Mon Sep 17 00:00:00 2001
From: zipg <geeky@foxmail.com>
Date: Mon, 3 Aug 2026 19:18:52 +0800
Subject: [PATCH] fix(agents): honor download source for update checks

---
 apps/desktop/src/App.vue                      |  5 +-
 crates/dbx-core/src/lib.rs                    | 14 ++----
 crates/dbx-mcp/src/transport.rs               | 46 +++++++++----------
 .../app-tests/driverDownloadSource.test.ts    |  3 ++
 4 files changed, 31 insertions(+), 37 deletions(-)

diff --git a/apps/desktop/src/App.vue b/apps/desktop/src/App.vue
index e48f5bfbaa..4aeff47a6a 100644
--- a/apps/desktop/src/App.vue
+++ b/apps/desktop/src/App.vue
@@ -1,7 +1,6 @@
 <script setup lang="ts">
 import { ref, computed, watch, onMounted, onUnmounted, nextTick, defineAsyncComponent } from "vue";
 import { useI18n } from "vue-i18n";
-import { invoke } from "@tauri-apps/api/core";
 import { ChevronsRight } from "@lucide/vue";
 import { TooltipProvider } from "@/components/ui/tooltip";
 import AppToolbar from "@/components/layout/AppToolbar.vue";
@@ -86,7 +85,7 @@ import { buildAppendedEditorSql } from "@/lib/ai/aiSqlAppend";
 import { assessProductionSql } from "@/lib/database/productionSafety";
 import { executeWithProductionSqlGuard } from "@/lib/database/productionExecutionGuard";
 import { buildHistoryAiAnalysisPrompt } from "@/lib/history/historyAiAnalysis";
-import { countAvailableAgentDriverUpdates, type AgentDriverUpdateBadgeState } from "@/lib/connection/agentDriverUpdateBadge";
+import { countAvailableAgentDriverUpdates } from "@/lib/connection/agentDriverUpdateBadge";
 import type { DriverStoreFocus } from "@/lib/connection/agentDriverInstallHint";
 import { safeLocalStorageGet, safeLocalStorageSet } from "@/lib/backend/safeStorage";
 import { apiUrl, webPath } from "@/lib/common/webPath";
@@ -245,7 +244,7 @@ function updateAgentDriverUpdateCount(count: number) {
 async function refreshAgentDriverUpdateCount() {
   if (!isDesktop || !settingsStore.editorSettings.updateNotificationsEnabled) return;
   try {
-    const drivers = await invoke<AgentDriverUpdateBadgeState[]>("list_installed_agents");
+    const drivers = await api.listInstalledAgents();
     if (!settingsStore.editorSettings.updateNotificationsEnabled) return;
     updateAgentDriverUpdateCount(countAvailableAgentDriverUpdates(drivers));
   } catch {
diff --git a/crates/dbx-core/src/lib.rs b/crates/dbx-core/src/lib.rs
index d2d1d470a5..da9e309d6d 100644
--- a/crates/dbx-core/src/lib.rs
+++ b/crates/dbx-core/src/lib.rs
@@ -106,7 +106,6 @@ impl DownloadSource {
         match self {
             Self::Official => Ok(download_candidate_urls(github_url, r2_path)),
             Self::Cnb => Ok(mirror_download_candidate_urls(
-                github_url,
                 r2_path,
                 rewrite_github_release_url(github_url, CNB_RELEASE_DOWNLOAD_PREFIX)?,
             )),
@@ -114,14 +113,9 @@ impl DownloadSource {
     }
 }
 
-fn mirror_download_candidate_urls(github_url: &str, r2_path: &str, mirror_url: String) -> Vec<String> {
+fn mirror_download_candidate_urls(r2_path: &str, mirror_url: String) -> Vec<String> {
     let r2_url = format!("{R2_CDN_BASE}{r2_path}");
-    // Mutable mirror aliases can lag even when versioned release assets are healthy.
-    if github_url.ends_with("/agents-latest/agent-registry.json") {
-        vec![r2_url, mirror_url]
-    } else {
-        vec![mirror_url, r2_url]
-    }
+    vec![mirror_url, r2_url]
 }
 
 fn rewrite_github_release_url(url: &str, target_prefix: &str) -> Result<String, String> {
@@ -202,13 +196,13 @@ mod tests {
     }
 
     #[test]
-    fn mirror_download_candidates_prefer_stable_registry_metadata() {
+    fn mirror_download_candidates_prefer_selected_source() {
         let github_url = "https://github.com/t8y2/dbx/releases/download/agents-latest/agent-registry.json";
         assert_eq!(
             DownloadSource::Cnb.download_candidate_urls(github_url, "agents/agent-registry.json").unwrap(),
             vec![
-                "https://dl.dbxio.com/agents/agent-registry.json",
                 "https://cnb.cool/dbxio.com/dbx/-/releases/download/agents-latest/agent-registry.json",
+                "https://dl.dbxio.com/agents/agent-registry.json",
             ]
         );
     }
diff --git a/crates/dbx-mcp/src/transport.rs b/crates/dbx-mcp/src/transport.rs
index 6e825ce4bf..669bc032d8 100644
--- a/crates/dbx-mcp/src/transport.rs
+++ b/crates/dbx-mcp/src/transport.rs
@@ -32,31 +32,29 @@ where
         self.inner.send(item)
     }
 
-    fn receive(&mut self) -> impl Future<Output = Option<RxJsonRpcMessage<RoleServer>>> + Send {
-        async move {
-            loop {
-                let message = self.inner.receive().await?;
-                let discovery_request_id = match &message {
-                    ClientJsonRpcMessage::Request(request)
-                        if matches!(
-                            &request.request,
-                            ClientRequest::CustomRequest(custom) if custom.method == "server/discover"
-                        ) =>
-                    {
-                        Some(request.id.clone())
-                    }
-                    _ => None,
-                };
-
-                let Some(request_id) = discovery_request_id else {
-                    return Some(message);
-                };
-
-                // Reject discovery without closing so new clients can fall back to the legacy initialize flow.
-                let error = ErrorData::new(ErrorCode::METHOD_NOT_FOUND, "Method not found", None);
-                if self.inner.send(ServerJsonRpcMessage::error(error, Some(request_id))).await.is_err() {
-                    return None;
+    async fn receive(&mut self) -> Option<RxJsonRpcMessage<RoleServer>> {
+        loop {
+            let message = self.inner.receive().await?;
+            let discovery_request_id = match &message {
+                ClientJsonRpcMessage::Request(request)
+                    if matches!(
+                        &request.request,
+                        ClientRequest::CustomRequest(custom) if custom.method == "server/discover"
+                    ) =>
+                {
+                    Some(request.id.clone())
                 }
+                _ => None,
+            };
+
+            let Some(request_id) = discovery_request_id else {
+                return Some(message);
+            };
+
+            // Reject discovery without closing so new clients can fall back to the legacy initialize flow.
+            let error = ErrorData::new(ErrorCode::METHOD_NOT_FOUND, "Method not found", None);
+            if self.inner.send(ServerJsonRpcMessage::error(error, Some(request_id))).await.is_err() {
+                return None;
             }
         }
     }
diff --git a/packages/app-tests/driverDownloadSource.test.ts b/packages/app-tests/driverDownloadSource.test.ts
index 4c6e3b5497..d4f38b653e 100644
--- a/packages/app-tests/driverDownloadSource.test.ts
+++ b/packages/app-tests/driverDownloadSource.test.ts
@@ -8,6 +8,7 @@ function source(relativePath: string): string {
 }
 
 test("shares the configured download source with agent driver management", () => {
+  const app = source("apps/desktop/src/App.vue");
   const driverStore = source("apps/desktop/src/components/config/DriverStoreDialog.vue");
   const backendApi = source("apps/desktop/src/lib/backend/api.ts");
 
@@ -20,4 +21,6 @@ test("shares the configured download source with agent driver management", () =>
   assert.match(backendApi, /backend\.installAgent\(dbType, useSettingsStore\(\)\.editorSettings\.updateDownloadSource(?:, operationId)?\)/);
   assert.match(backendApi, /backend\.upgradeAllAgents\(useSettingsStore\(\)\.editorSettings\.updateDownloadSource(?:, operationId)?\)/);
   assert.match(backendApi, /backend\.reinstallJre\(jreKey, useSettingsStore\(\)\.editorSettings\.updateDownloadSource(?:, operationId)?\)/);
+  assert.match(app, /const drivers = await api\.listInstalledAgents\(\)/);
+  assert.doesNotMatch(app, /invoke<[^>]+>\("list_installed_agents"\)/);
 });