fix(plugins): keep borrowed host connections on plugin tab close

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();
   });
 });

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<string, unknown>; 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<string>();
   /**
-   * 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<QueryTab>) {
-    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<string, Set<string>>();
+    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<string>();
+      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<string, unknown>, performance.now() - commandStartedAt))));
                 mongoEditTarget = undefined;
                 queryExecutionLog("info", "mongo-collection-stats:done", {