feat(python-setup): open the compute picker inline when no compute is attached

packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts

@@ -3,11 +3,57 @@ import {ApiClient} from "@databricks/sdk-experimental";
 import {Cluster} from "../sdk-extensions";
 import assert from "assert";
 import {mock} from "ts-mockito";
+import {QuickPickItem, window} from "vscode";
 import {
+    ConnectionCommands,
     formatClusterState,
     formatQuickPickClusterDetails,
 } from "./ConnectionCommands";
 
+/**
+ * Scriptable stand-in for the compute QuickPick. The test drives
+ * {@link accept}/{@link hide} explicitly, since the command registers its
+ * handlers only after `show()`.
+ */
+class FakeComputeQuickPick {
+    title?: string;
+    keepScrollPosition = false;
+    busy = false;
+    canSelectMany = false;
+    items: readonly QuickPickItem[] = [];
+    selectedItems: readonly QuickPickItem[] = [];
+    disposed = false;
+    private readonly acceptCbs: Array<() => void | Promise<void>> = [];
+    private readonly hideCbs: Array<() => void> = [];
+
+    onDidAccept(cb: () => void | Promise<void>) {
+        this.acceptCbs.push(cb);
+        return {dispose() {}};
+    }
+    onDidHide(cb: () => void) {
+        this.hideCbs.push(cb);
+        return {dispose() {}};
+    }
+    show() {}
+    dispose() {
+        this.disposed = true;
+    }
+
+    /** Simulate the user picking `selected` and pressing Enter. */
+    async accept(selected: readonly QuickPickItem[]) {
+        this.selectedItems = selected;
+        for (const cb of this.acceptCbs) {
+            await cb();
+        }
+    }
+    /** Simulate the user dismissing the picker (Escape). */
+    async hide() {
+        for (const cb of this.hideCbs) {
+            cb();
+        }
+    }
+}
+
 describe(__filename, () => {
     it("attach cluster quickpick: correctly format cluster details", () => {
         const clusterDetails = formatQuickPickClusterDetails(
@@ -34,4 +80,157 @@ describe(__filename, () => {
         assert.equal(formatClusterState("ERROR"), "Error");
         assert.equal(formatClusterState("UNKNOWN"), "Unknown");
     });
+
+    describe("attachClusterQuickPick return contract", () => {
+        let originalCreateQuickPick: typeof window.createQuickPick;
+        let fakePick: FakeComputeQuickPick;
+        let attachCalls: string[];
+        let commands: ConnectionCommands;
+
+        beforeEach(() => {
+            originalCreateQuickPick = window.createQuickPick;
+            fakePick = new FakeComputeQuickPick();
+            (window as unknown as {createQuickPick: unknown}).createQuickPick =
+                () => fakePick;
+
+            attachCalls = [];
+            const connectionManager = {
+                workspaceClient: {},
+                databricksWorkspace: {userName: "me"},
+                attachCluster: async (id: string) => {
+                    attachCalls.push(id);
+                },
+                enableServerless: async () => {},
+            };
+            const clusterModel = {
+                refresh() {},
+                onDidChange() {
+                    return {dispose() {}};
+                },
+                roots: [],
+            };
+            commands = new ConnectionCommands(
+                {} as never,
+                connectionManager as never,
+                clusterModel as never,
+                {} as never,
+                {} as never,
+                {} as never
+            );
+        });
+
+        afterEach(() => {
+            (window as unknown as {createQuickPick: unknown}).createQuickPick =
+                originalCreateQuickPick;
+        });
+
+        const clusterItem = (id: string) =>
+            ({
+                label: `cluster ${id}`,
+                cluster: {id},
+            }) as unknown as QuickPickItem;
+
+        it("resolves to the attached cluster and attaches it exactly once", async () => {
+            const resultP = commands.attachClusterQuickPickCommand()();
+            await fakePick.accept([clusterItem("c1")]);
+
+            assert.deepEqual(await resultP, {
+                kind: "cluster",
+                clusterId: "c1",
+            });
+            assert.deepEqual(attachCalls, ["c1"]);
+        });
+
+        it("resolves to undefined when the picker is dismissed", async () => {
+            const resultP = commands.attachClusterQuickPickCommand()();
+            await fakePick.hide();
+
+            assert.equal(await resultP, undefined);
+            assert.deepEqual(attachCalls, []);
+        });
+
+        it("guards re-entry: a second Enter neither re-attaches nor changes the result", async () => {
+            // Without the `settled` guard a double Enter would attach twice.
+            const resultP = commands.attachClusterQuickPickCommand()();
+            await fakePick.accept([clusterItem("c1")]);
+            await fakePick.accept([clusterItem("c2")]);
+
+            assert.deepEqual(await resultP, {
+                kind: "cluster",
+                clusterId: "c1",
+            });
+            assert.deepEqual(attachCalls, ["c1"]);
+        });
+
+        it("resolves to undefined when accepted with nothing highlighted", async () => {
+            const resultP = commands.attachClusterQuickPickCommand()();
+            await fakePick.accept([]);
+
+            assert.equal(await resultP, undefined);
+            assert.deepEqual(attachCalls, []);
+        });
+
+        it("still returns the chosen cluster when the attach silently fails (best-effort)", async () => {
+            // attachCluster is @onError(throw:false): a failed write resolves
+            // without throwing, so we still return the picked target by design.
+            const attach = async () => {};
+            const cmds = new ConnectionCommands(
+                {} as never,
+                {
+                    workspaceClient: {},
+                    databricksWorkspace: {userName: "me"},
+                    attachCluster: attach,
+                    enableServerless: async () => {},
+                } as never,
+                {
+                    refresh() {},
+                    onDidChange() {
+                        return {dispose() {}};
+                    },
+                    roots: [],
+                } as never,
+                {} as never,
+                {} as never,
+                {} as never
+            );
+
+            const resultP = cmds.attachClusterQuickPickCommand()();
+            await fakePick.accept([clusterItem("c1")]);
+
+            assert.deepEqual(await resultP, {
+                kind: "cluster",
+                clusterId: "c1",
+            });
+        });
+
+        it("resolves to undefined if an unexpected error is thrown during selection", async () => {
+            // A throw from any step must settle the Promise, not hang.
+            const cmds = new ConnectionCommands(
+                {} as never,
+                {
+                    workspaceClient: {},
+                    databricksWorkspace: {userName: "me"},
+                    attachCluster: async () => {
+                        throw new Error("unexpected");
+                    },
+                    enableServerless: async () => {},
+                } as never,
+                {
+                    refresh() {},
+                    onDidChange() {
+                        return {dispose() {}};
+                    },
+                    roots: [],
+                } as never,
+                {} as never,
+                {} as never,
+                {} as never
+            );
+
+            const resultP = cmds.attachClusterQuickPickCommand()();
+            await fakePick.accept([clusterItem("c1")]);
+
+            assert.equal(await resultP, undefined);
+        });
+    });
 });

packages/databricks-vscode/src/configuration/ConnectionCommands.ts

@@ -1,5 +1,5 @@
 import {Cluster} from "../sdk-extensions";
-import {compute} from "@databricks/sdk-experimental";
+import {compute, logging} from "@databricks/sdk-experimental";
 import {
     Disposable,
     QuickPickItem,
@@ -30,7 +30,19 @@ import {
 } from "../python-setup/utils/serverlessVersionResolver";
 import {pickServerlessVersion} from "../python-setup/utils/serverlessVersionPicker";
 import {collectServerlessVersionObservations} from "../python-setup/utils/serverlessVersionObservations";
+import type {SetupCompute} from "../python-setup/controllers/PythonSetupEnvironmentSetup";
 import {WorkspaceFolderManager} from "../vscode-objs/WorkspaceFolderManager";
+import {Loggers} from "../logger";
+
+// eslint-disable-next-line @typescript-eslint/naming-convention
+const {NamedLogger} = logging;
+
+/**
+ * A compute target picked in the QuickPick. Aliased to the `setup-local` compute
+ * shape so a drift is a compile error, not a silent mismatch through
+ * `executeCommand`'s untyped generic. Serverless is always version-complete.
+ */
+export type SelectedCompute = SetupCompute;
 
 function formatQuickPickClusterSize(sizeInMB: number): string {
     if (sizeInMB > 1024) {
@@ -143,8 +155,14 @@ export class ConnectionCommands implements Disposable {
         };
     }
 
+    /**
+     * Resolves once the picker closes to the attached compute, or `undefined` on
+     * dismissal / "Create New Cluster". The return lets python-setup use the
+     * selection directly rather than re-reading the connection manager (whose
+     * cluster attach is async and would race). Other callers ignore it.
+     */
     attachClusterQuickPickCommand() {
-        return async (title?: string) => {
+        return async (title?: string): Promise<SelectedCompute | undefined> => {
             const workspaceClient = this.connectionManager.workspaceClient;
             const me = this.connectionManager.databricksWorkspace?.userName;
             if (!workspaceClient || !me) {
@@ -204,33 +222,73 @@ export class ConnectionCommands implements Disposable {
             refreshQuickPickItems();
             quickPick.show();
 
-            quickPick.onDidAccept(async () => {
-                const selectedItem = quickPick.selectedItems[0];
-                if ("cluster" in selectedItem) {
-                    const cluster = selectedItem.cluster;
-                    await this.connectionManager.attachCluster(cluster.id);
-                } else if (selectedItem.label === "$(cloud) Serverless") {
-                    // Dispose the compute QuickPick before opening the version
-                    // sub-picker so they don't stack visually.
-                    disposables.forEach((d) => d.dispose());
-                    await this.selectServerless();
-                    return;
-                } else {
-                    await UrlUtils.openExternal(
-                        `${
-                            (
-                                await this.connectionManager.workspaceClient
-                                    ?.apiClient?.host
-                            )?.href ?? ""
-                        }#create/cluster`
-                    );
-                }
-                disposables.forEach((d) => d.dispose());
-            });
+            // `settled` guards a repeated Enter and stops the hide handler
+            // (which also fires when accept disposes the picker) from resolving
+            // before the accept branch's awaits finish.
+            return await new Promise<SelectedCompute | undefined>((resolve) => {
+                let settled = false;
+                quickPick.onDidAccept(async () => {
+                    if (settled) {
+                        return;
+                    }
+                    settled = true;
+                    let selected: SelectedCompute | undefined;
+                    try {
+                        const selectedItem = quickPick.selectedItems[0];
+                        if (selectedItem === undefined) {
+                            // Accepted with nothing highlighted.
+                        } else if ("cluster" in selectedItem) {
+                            const cluster = selectedItem.cluster;
+                            // Best-effort: attachCluster is @onError(throw:false),
+                            // so a failed write pops its own error but doesn't
+                            // throw. We still return the chosen target -- it's all
+                            // setup-local needs and any failure is already shown.
+                            await this.connectionManager.attachCluster(
+                                cluster.id
+                            );
+                            selected = {kind: "cluster", clusterId: cluster.id};
+                        } else if (
+                            selectedItem.label === "$(cloud) Serverless"
+                        ) {
+                            // Dispose the compute QuickPick before opening the
+                            // version sub-picker so they don't stack visually.
+                            disposables.forEach((d) => d.dispose());
+                            const version = await this.selectServerless();
+                            if (version !== undefined) {
+                                selected = {kind: "serverless", version};
+                            }
+                        } else {
+                            await UrlUtils.openExternal(
+                                `${
+                                    (
+                                        await this.connectionManager
+                                            .workspaceClient?.apiClient?.host
+                                    )?.href ?? ""
+                                }#create/cluster`
+                            );
+                        }
+                    } catch (e) {
+                        // Defense-in-depth for an unexpected throw (attach/enable
+                        // are @onError and don't throw): keep it off the unhandled
+                        // path and still settle. `selected` stays undefined.
+                        NamedLogger.getOrCreate(Loggers.Extension).error(
+                            "Compute picker selection failed",
+                            e
+                        );
+                    } finally {
+                        disposables.forEach((d) => d.dispose());
+                        resolve(selected);
+                    }
+                });
 
-            quickPick.onDidHide(() => {
-                disposables.forEach((d) => d.dispose());
-                quickPick.dispose();
+                quickPick.onDidHide(() => {
+                    disposables.forEach((d) => d.dispose());
+                    quickPick.dispose();
+                    if (!settled) {
+                        settled = true;
+                        resolve(undefined);
+                    }
+                });
             });
         };
     }
@@ -242,18 +300,22 @@ export class ConnectionCommands implements Disposable {
      * selection, so setup need not re-prompt. If they dismiss the version
      * picker, no compute change is made. With the feature off this is the
      * plain, unchanged serverless enable.
+     *
+     * Returns the confirmed version, or `undefined` if the picker was dismissed
+     * or the feature is off (serverless enabled but version-less).
      */
-    private async selectServerless() {
+    private async selectServerless(): Promise<string | undefined> {
         if (!isPythonSetupEnabled()) {
             await this.connectionManager.enableServerless();
-            return;
+            return undefined;
         }
         const version = await this.pickServerlessVersion();
         if (version === undefined) {
             // User dismissed the version picker -- don't switch compute.
-            return;
+            return undefined;
         }
         await this.connectionManager.enableServerless(version);
+        return version;
     }
 
     /**