From c7660b20d0b5483aa669f0bf308b21ee38c9294b Mon Sep 17 00:00:00 2001 From: jearhe <0verme7753@gmail.com> Date: Mon, 21 Sep 2026 11:36:22 +0800 Subject: [PATCH] fix(plugins): keep borrowed host connections on plugin tab close --- ...ueryStore.pluginTabCloseDisconnect.spec.ts | 130 ++++++++++++++++-- apps/desktop/src/stores/queryStore.ts | 37 +++-- 2 files changed, 146 insertions(+), 21 deletions(-) diff --git a/apps/desktop/src/stores/__tests__/queryStore.pluginTabCloseDisconnect.spec.ts b/apps/desktop/src/stores/__tests__/queryStore.pluginTabCloseDisconnect.spec.ts index 0bae492ed9..370bbd04b3 100644 --- a/apps/desktop/src/stores/__tests__/queryStore.pluginTabCloseDisconnect.spec.ts +++ b/apps/desktop/src/stores/__tests__/queryStore.pluginTabCloseDisconnect.spec.ts @@ -2,8 +2,37 @@ import { createPinia, setActivePinia } from "pinia"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { useConnectionStore } from "@/stores/connectionStore"; import { useQueryStore } from "@/stores/queryStore"; +import type { ConnectionConfig } from "@/types/database"; -describe("queryStore plugin tab close disconnects the plugin connection", () => { +function postgresConnection(id: string): ConnectionConfig { + return { + id, + name: "PostgreSQL", + db_type: "postgres", + host: "127.0.0.1", + port: 5432, + username: "postgres", + password: "", + database: "app", + }; +} + +function pluginConnection(id: string, pluginId: string): ConnectionConfig { + return { + ...postgresConnection(id), + name: "Plugin connection", + db_type: "plugin", + plugin_id: pluginId, + plugin_connection_provider: "connection", + plugin_connection_type: "ssh", + }; +} + +async function settlePluginTabClose() { + await new Promise((resolve) => setTimeout(resolve, 20)); +} + +describe("queryStore plugin tab close connection lifecycle", () => { beforeEach(() => { vi.resetModules(); vi.unstubAllGlobals(); @@ -20,41 +49,116 @@ describe("queryStore plugin tab close disconnects the plugin connection", () => setActivePinia(createPinia()); }); - it("disconnects the plugin connection when its last plugin tab closes", async () => { + it("keeps a host-owned PostgreSQL connection and SQL tab when its plugin tab closes", async () => { const queryStore = useQueryStore(); const connectionStore = useConnectionStore(); + connectionStore.connections = [postgresConnection("conn-host")]; const disconnectSpy = vi.spyOn(connectionStore, "disconnect").mockResolvedValue(); - const id = queryStore.openPluginWorkbench("io.dbx.ssh", "workbench", { connectionId: "conn-1", context: { connectionId: "conn-1" } }); - queryStore.closeTab(id); + const sqlTabId = queryStore.createTab("conn-host", "app", "Query 1", "query"); + const pluginTabId = queryStore.openPluginWorkbench("io.dbx.ssh", "workbench", { + connectionId: "conn-host", + context: { connectionId: "conn-host" }, + }); - await vi.waitFor(() => expect(disconnectSpy).toHaveBeenCalledTimes(1)); - expect(disconnectSpy).toHaveBeenCalledWith("conn-1"); + queryStore.closeTab(pluginTabId); + await settlePluginTabClose(); + + expect(disconnectSpy).not.toHaveBeenCalled(); + expect(queryStore.tabs.some((tab) => tab.id === sqlTabId)).toBe(true); }); - it("keeps the plugin connection while another plugin tab for it stays open", async () => { + it("never disconnects a host-owned connection after its last plugin consumer closes", async () => { const queryStore = useQueryStore(); const connectionStore = useConnectionStore(); + connectionStore.connections = [postgresConnection("conn-host")]; const disconnectSpy = vi.spyOn(connectionStore, "disconnect").mockResolvedValue(); - // Workbench + filesystem tab riding the same connection (e.g. SSH terminal + SFTP). - const workbenchId = queryStore.openPluginWorkbench("io.dbx.ssh", "workbench", { connectionId: "conn-1", context: { connectionId: "conn-1" } }); - queryStore.openPluginFilesystem("io.dbx.ssh", "fs", { connectionId: "conn-1", rootUri: "sftp:/" }); + const workbenchId = queryStore.openPluginWorkbench("io.dbx.ssh", "workbench", { + connectionId: "conn-host", + context: { connectionId: "conn-host" }, + }); + const filesystemId = queryStore.openPluginFilesystem("io.dbx.ssh", "fs", { + connectionId: "conn-host", + rootUri: "sftp:/", + }); queryStore.closeTab(workbenchId); - await new Promise((resolve) => setTimeout(resolve, 20)); + await settlePluginTabClose(); + expect(disconnectSpy).not.toHaveBeenCalled(); + + queryStore.closeTab(filesystemId); + await settlePluginTabClose(); expect(disconnectSpy).not.toHaveBeenCalled(); }); - it("does not disconnect non-plugin connections when a regular tab closes", async () => { + it("disconnects a plugin-owned connection when its last owner tab closes", async () => { const queryStore = useQueryStore(); const connectionStore = useConnectionStore(); + connectionStore.connections = [pluginConnection("conn-plugin", "io.dbx.ssh")]; + const disconnectSpy = vi.spyOn(connectionStore, "disconnect").mockResolvedValue(); + + const pluginTabId = queryStore.openPluginWorkbench("io.dbx.ssh", "workbench", { + connectionId: "conn-plugin", + context: { connectionId: "conn-plugin" }, + }); + + queryStore.closeTab(pluginTabId); + await vi.waitFor(() => expect(disconnectSpy).toHaveBeenCalledTimes(1)); + expect(disconnectSpy).toHaveBeenCalledWith("conn-plugin"); + }); + + it("waits for the last plugin-owned workbench and filesystem tab before disconnecting", async () => { + const queryStore = useQueryStore(); + const connectionStore = useConnectionStore(); + connectionStore.connections = [pluginConnection("conn-plugin", "io.dbx.ssh")]; + const disconnectSpy = vi.spyOn(connectionStore, "disconnect").mockResolvedValue(); + + const workbenchId = queryStore.openPluginWorkbench("io.dbx.ssh", "workbench", { + connectionId: "conn-plugin", + context: { connectionId: "conn-plugin" }, + }); + const filesystemId = queryStore.openPluginFilesystem("io.dbx.ssh", "fs", { + connectionId: "conn-plugin", + rootUri: "sftp:/", + }); + + queryStore.closeTab(workbenchId); + await settlePluginTabClose(); + expect(disconnectSpy).not.toHaveBeenCalled(); + + queryStore.closeTab(filesystemId); + await vi.waitFor(() => expect(disconnectSpy).toHaveBeenCalledTimes(1)); + expect(disconnectSpy).toHaveBeenCalledWith("conn-plugin"); + }); + + it("does not let another plugin disconnect a plugin-owned connection it only borrows", async () => { + const queryStore = useQueryStore(); + const connectionStore = useConnectionStore(); + connectionStore.connections = [pluginConnection("conn-plugin", "io.dbx.owner")]; + const disconnectSpy = vi.spyOn(connectionStore, "disconnect").mockResolvedValue(); + + const borrowedTabId = queryStore.openPluginWorkbench("io.dbx.borrower", "workbench", { + connectionId: "conn-plugin", + context: { connectionId: "conn-plugin" }, + }); + + queryStore.closeTab(borrowedTabId); + await settlePluginTabClose(); + + expect(disconnectSpy).not.toHaveBeenCalled(); + }); + + it("does not disconnect a non-plugin connection when a regular tab closes", async () => { + const queryStore = useQueryStore(); + const connectionStore = useConnectionStore(); + connectionStore.connections = [postgresConnection("pg-1")]; const disconnectSpy = vi.spyOn(connectionStore, "disconnect").mockResolvedValue(); const id = queryStore.createTab("pg-1", "app", "Query 1", "query"); queryStore.closeTab(id); - await new Promise((resolve) => setTimeout(resolve, 20)); + await settlePluginTabClose(); expect(disconnectSpy).not.toHaveBeenCalled(); }); }); diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index a90d1032ba..7f7c4e341e 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -3494,6 +3494,12 @@ export const useQueryStore = defineStore("query", () => { return typeof contextConnectionId === "string" ? contextConnectionId : ""; } + function pluginTabPluginId(tab: QueryTab): string | undefined { + if (tab.mode === "plugin-workbench") return tab.pluginWorkbench?.pluginId; + if (tab.mode === "plugin-filesystem") return tab.pluginFilesystem?.pluginId; + return undefined; + } + function openPluginWorkbench(pluginId: string, contributionId: string, options: { title?: string; connectionId?: string; database?: string; context?: Record; forceNew?: boolean } = {}) { const contextConnectionId = typeof options.context?.connectionId === "string" ? options.context.connectionId : ""; const connectionId = options.connectionId || contextConnectionId; @@ -3962,15 +3968,25 @@ export const useQueryStore = defineStore("query", () => { const pluginReleaseInFlight = new Set(); /** - * Closing the last plugin tab (workbench or filesystem) bound to a plugin - * connection ends that connection: the sidecar session has no remaining - * owner, so disconnecting tears the pool down and drops the sidebar's - * "connected" dot. Fire-and-forget; concurrent releases for the same - * connection coalesce, and a reopen racing the release wins. + * Closing the last plugin tab only releases a connection explicitly owned by + * the plugin that closed it. Workbench/filesystem tabs can borrow a Host-owned + * connection, so a connectionId reference alone is never enough to disconnect. + * Fire-and-forget; concurrent releases for the same connection coalesce, and + * a reopen racing the release wins. */ function releasePluginConnectionsAfterClose(closedTabs: ReadonlyArray) { - const connectionIds = new Set(closedTabs.filter((tab) => (tab.mode === "plugin-workbench" || tab.mode === "plugin-filesystem") && tab.connectionId).map((tab) => tab.connectionId)); - for (const connectionId of connectionIds) { + const closedPluginTabsByConnection = new Map>(); + for (const tab of closedTabs) { + if (tab.mode !== "plugin-workbench" && tab.mode !== "plugin-filesystem") continue; + const connectionId = pluginTabConnectionId(tab); + const pluginId = pluginTabPluginId(tab); + if (!connectionId || !pluginId) continue; + const pluginIds = closedPluginTabsByConnection.get(connectionId) ?? new Set(); + pluginIds.add(pluginId); + closedPluginTabsByConnection.set(connectionId, pluginIds); + } + + for (const [connectionId, closedPluginIds] of closedPluginTabsByConnection) { if (pluginReleaseInFlight.has(connectionId)) continue; pluginReleaseInFlight.add(connectionId); void (async () => { @@ -3978,8 +3994,12 @@ export const useQueryStore = defineStore("query", () => { // Yield one microtask so a same-tick reopen registers its // replacement plugin tab before we tear the connection down. await Promise.resolve(); - if (tabs.value.some((tab) => (tab.mode === "plugin-workbench" || tab.mode === "plugin-filesystem") && tab.connectionId === connectionId)) return; + if (tabs.value.some((tab) => (tab.mode === "plugin-workbench" || tab.mode === "plugin-filesystem") && pluginTabConnectionId(tab) === connectionId)) return; const connectionStore = useConnectionStore(); + const config = connectionStore.getConfig(connectionId); + // A connection reference is not ownership: only a plugin-backed + // connection closed by its owning plugin may be auto-released. + if (config?.db_type !== "plugin" || !config.plugin_id || !closedPluginIds.has(config.plugin_id)) return; // A user-initiated disconnect that closed these tabs already owns // the teardown — don't stack a second one on top of it. if (connectionStore.hasDisconnectInFlight(connectionId)) return; @@ -6856,6 +6876,7 @@ export const useQueryStore = defineStore("query", () => { database: currentDatabase, }); const stats = await api.mongoCollectionStats(executionConnectionId, currentDatabase, mongoCommand.collection, mongoCommand.scale, executionId); + // SAFETY: The backend returns collection statistics as a JSON object; the API type is broader than the converter's record-shaped input. allResults.push(markQueryResultRowsRaw(annotateMongoResult(mongoCollectionStatsToQueryResult(mongoCommand.metric, stats as unknown as Record, performance.now() - commandStartedAt)))); mongoEditTarget = undefined; queryExecutionLog("info", "mongo-collection-stats:done", {