From 0647e898440d188dacb87255da2e3aeafa8dda26 Mon Sep 17 00:00:00 2001 From: "@rugpanov" Date: Thu, 13 Aug 2026 17:15:13 +0200 Subject: [PATCH 1/5] feat(python-setup): open the compute picker inline when none is attached Why: Clicking "Set up Python environment" with no compute attached dead-ended on a "select a cluster or serverless compute first" warning, forcing the user to leave the flow, attach compute, and come back. A serverless session missing only its version already prompts inline and continues; the no-compute case should be just as direct. What: - resolveCompute (pythonSetupDeps): when nothing is attached, open the existing compute picker inline, then re-read the attachment and continue setup against whatever the user selected. Dismissing the picker leaves the attachment unchanged, so the flow still falls through to the existing guidance message + no-compute telemetry -- nothing regresses, the picker is only an added opportunity. - attachClusterQuickPick command (ConnectionCommands): its returned promise now resolves only once the picker has fully closed (after any attach/enable it triggers, or immediately on dismissal), so a caller can await the outcome and re-read state. Existing fire-and-forget callers ignore the result unchanged. - extension.ts: wire promptSelectCompute to the compute-picker command. Verification: - yarn build; unit suite green (462 passing) incl. new resolveCompute cases: picker opened when none attached, runs against the attached compute, not opened when a cluster is already attached, version prompted when the picker leaves serverless version-less. - yarn fix / test:lint clean. Co-authored-by: Isaac --- .../src/configuration/ConnectionCommands.ts | 81 +++++++++----- packages/databricks-vscode/src/extension.ts | 11 ++ .../controllers/pythonSetupDeps.test.ts | 100 ++++++++++++++++-- .../controllers/pythonSetupDeps.ts | 20 +++- 4 files changed, 178 insertions(+), 34 deletions(-) diff --git a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts b/packages/databricks-vscode/src/configuration/ConnectionCommands.ts index b930d2483..ef8a1055d 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts +++ b/packages/databricks-vscode/src/configuration/ConnectionCommands.ts @@ -143,8 +143,15 @@ export class ConnectionCommands implements Disposable { }; } + /** + * The returned handler resolves only once the picker has fully closed -- + * after any attach/enable it triggers has completed, or immediately when the + * user dismisses it. Callers that just open the picker can ignore the result + * (they always have); the resolution lets a caller that needs the outcome + * (python-setup, which re-reads the attached compute afterward) await it. + */ attachClusterQuickPickCommand() { - return async (title?: string) => { + return async (title?: string): Promise => { const workspaceClient = this.connectionManager.workspaceClient; const me = this.connectionManager.databricksWorkspace?.userName; if (!workspaceClient || !me) { @@ -204,33 +211,53 @@ 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()); - }); + // Resolve only once the picker's work is done: on accept, after the + // attach/enable (or browser hand-off) it triggers; on a plain hide, + // right away. `accepted` keeps the hide handler -- which also fires + // when accept disposes the picker -- from resolving early, before the + // accept branch's awaits have settled. + await new Promise((resolve) => { + let accepted = false; + quickPick.onDidAccept(async () => { + accepted = true; + try { + 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()); + } finally { + resolve(); + } + }); - quickPick.onDidHide(() => { - disposables.forEach((d) => d.dispose()); - quickPick.dispose(); + quickPick.onDidHide(() => { + disposables.forEach((d) => d.dispose()); + quickPick.dispose(); + if (!accepted) { + resolve(); + } + }); }); }; } diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index c34f05d1a..e69690f6d 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -1005,6 +1005,17 @@ export async function activate( // setup flow. persistServerlessVersion: (version) => connectionManager.enableServerless(version), + // Reuses the existing compute picker. When the feature is opted in, + // its serverless branch also records the environment version, so a + // serverless selection made here comes back version-complete and + // setup need not re-prompt. Resolves once the picker closes; the + // setup flow then re-reads the attached compute. + promptSelectCompute: () => + Promise.resolve( + commands.executeCommand( + "databricks.connection.attachClusterQuickPick" + ) + ), setActiveInterpreter: async (interpreterPath, root) => { await pythonExtensionWrapper.api.environments.updateActiveEnvironmentPath( interpreterPath, diff --git a/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.test.ts b/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.test.ts index abc4eeb4d..93a6134d3 100644 --- a/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.test.ts +++ b/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.test.ts @@ -175,6 +175,7 @@ function makeWiring( }), promptServerlessVersion: async () => "4", persistServerlessVersion: async () => {}, + promptSelectCompute: async () => {}, setActiveInterpreter: async () => {}, persistSetupState: () => {}, log: {append: () => {}, show: () => {}}, @@ -254,21 +255,108 @@ describe("makePythonSetupDeps resolveCompute", () => { expect(prompted).to.equal(0); }); - it("never prompts when nothing is attached", async () => { - // Nothing selected is a real dead end, not a missing detail: the flow - // must guide the user, not open a version picker. - let prompted = 0; + it("offers the compute picker (not the version picker) when nothing is attached", async () => { + // Nothing selected is a missing target, not a missing detail: the flow + // opens the compute picker so the user can choose one -- it must never + // jump straight to the serverless version picker. + let versionPrompted = 0; + let pickerOpened = 0; const deps = makePythonSetupDeps( makeWiring({ + promptSelectCompute: async () => { + pickerOpened++; + // User dismisses the picker without attaching anything. + }, promptServerlessVersion: async () => { - prompted++; + versionPrompted++; return "4"; }, }) ); + // Dismissed with nothing attached -> still the dead-end `none`, which + // the orchestrator turns into the guidance message. expect(await deps.resolveCompute()).to.deep.equal({status: "none"}); - expect(prompted).to.equal(0); + expect(pickerOpened).to.equal(1); + expect(versionPrompted).to.equal(0); + }); + + it("runs against the compute the user attaches through the picker", async () => { + // The picker persists the selection through the connection manager, so + // resolveCompute re-reads the attachment after it closes. + let attached = { + serverless: false, + cluster: undefined as {id: string} | undefined, + serverlessVersion: undefined as string | undefined, + }; + const deps = makePythonSetupDeps( + makeWiring({ + attachedCompute: () => attached, + promptSelectCompute: async () => { + attached = { + serverless: false, + cluster: {id: "c1"}, + serverlessVersion: undefined, + }; + }, + }) + ); + + expect(await deps.resolveCompute()).to.deep.equal({ + status: "ok", + compute: {kind: "cluster", clusterId: "c1"}, + }); + }); + + it("does not open the compute picker when a cluster is already attached", async () => { + let pickerOpened = 0; + const deps = makePythonSetupDeps( + makeWiring({ + attachedCompute: () => ({ + serverless: false, + cluster: {id: "c1"}, + serverlessVersion: undefined, + }), + promptSelectCompute: async () => { + pickerOpened++; + }, + }) + ); + + expect(await deps.resolveCompute()).to.deep.equal({ + status: "ok", + compute: {kind: "cluster", clusterId: "c1"}, + }); + expect(pickerOpened).to.equal(0); + }); + + it("prompts for a version when the picker attaches version-less serverless", async () => { + // The picker can leave serverless selected without a version (a config + // that enables serverless without opening the version sub-picker). The + // flow then resolves the version rather than dead-ending. + let attached = { + serverless: false, + cluster: undefined as {id: string} | undefined, + serverlessVersion: undefined as string | undefined, + }; + const deps = makePythonSetupDeps( + makeWiring({ + attachedCompute: () => attached, + promptSelectCompute: async () => { + attached = { + serverless: true, + cluster: undefined, + serverlessVersion: undefined, + }; + }, + promptServerlessVersion: async () => "4", + }) + ); + + expect(await deps.resolveCompute()).to.deep.equal({ + status: "ok", + compute: {kind: "serverless", version: "4"}, + }); }); it("never prompts when serverless already has a version", async () => { diff --git a/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.ts b/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.ts index 2e294ab39..006df2aa9 100644 --- a/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.ts +++ b/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.ts @@ -137,6 +137,14 @@ export interface PythonSetupWiringDeps { * selection, so the next run does not ask again. */ persistServerlessVersion: (version: string) => Promise; + /** + * Open the compute picker so the user can attach a target when none is + * selected, resolving once it closes. The selection is persisted through + * the connection manager (a cluster attaches, serverless enables), so the + * caller re-reads {@link attachedCompute} afterward rather than taking a + * return value -- and re-reads `none` when the user dismissed it. + */ + promptSelectCompute: () => Promise; /** Point the MS Python extension at an interpreter path (project-scoped). */ setActiveInterpreter: (interpreterPath: string, root: Uri) => Promise; /** Persist the post-setup state (workspace-scoped) for drift detection. */ @@ -169,7 +177,17 @@ export function makePythonSetupDeps( projectRoot: wiring.projectRoot, isVisible, resolveCompute: async () => { - const resolution = resolveComputeFrom(wiring.attachedCompute()); + let resolution = resolveComputeFrom(wiring.attachedCompute()); + if (resolution.status === "none") { + // Nothing attached: open the compute picker inline so the user + // can choose a target, rather than dead-ending the CTA. The + // picker persists the selection through the connection manager, + // so re-read the attachment afterward -- it reflects whatever + // they chose, or stays `none` if they dismissed the picker (in + // which case the orchestrator shows the guidance message). + await wiring.promptSelectCompute(); + resolution = resolveComputeFrom(wiring.attachedCompute()); + } if (resolution.status !== "needsServerlessVersion") { return resolution; } From 82d3cc2985c44f3b04feedc5467b4ba6351c39c1 Mon Sep 17 00:00:00 2001 From: "@rugpanov" Date: Thu, 13 Aug 2026 17:35:19 +0200 Subject: [PATCH 2/5] fix(python-setup): use the picker's returned compute, not a racy re-read Why: Code review (Codex + a Claude reviewer, converging independently) caught that re-reading the attached compute right after the picker closed is deterministically stale for a cluster: ConnectionManager.attachCluster only writes the `clusterId` config, and the cluster object is rebuilt by a separate, network-gated, fire-and-forget config listener that the attach never awaits. So the immediate re-read saw `cluster: undefined` and setup dead-ended on the very message this feature removes. (The serverless path only worked by luck -- enableServerless sets its state synchronously.) The first cut's passing test masked this by mutating the mock synchronously. What: - attachClusterQuickPick command now RETURNS the chosen compute (SelectedCompute: cluster or version-complete serverless) or undefined on dismissal, so callers use the selection directly. selectServerless returns the confirmed version. Existing fire-and-forget callers ignore the new return value, unchanged. - resolveCompute (pythonSetupDeps) consumes that return value instead of re-reading connectionManager state, eliminating the race. - Harden the accept handler while here (review nits): guard a repeated Enter / empty selection, and catch+log so a failing attach can't leak an unhandled rejection or leave the picker open. - Tests rewritten to model the real contract: the picker returns the compute while attachedCompute stays `none`, proving the flow no longer depends on a re-read. Verification: - yarn build; full unit suite green (736 passing). yarn fix / test:lint clean. Co-authored-by: Isaac --- .../src/configuration/ConnectionCommands.ts | 88 ++++++++++++++----- packages/databricks-vscode/src/extension.ts | 13 +-- .../controllers/pythonSetupDeps.test.ts | 81 +++++++++-------- .../controllers/pythonSetupDeps.ts | 31 ++++--- 4 files changed, 130 insertions(+), 83 deletions(-) diff --git a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts b/packages/databricks-vscode/src/configuration/ConnectionCommands.ts index ef8a1055d..1689d0fbd 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts +++ b/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, @@ -32,6 +32,19 @@ import {pickServerlessVersion} from "../python-setup/utils/serverlessVersionPick import {collectServerlessVersionObservations} from "../python-setup/utils/serverlessVersionObservations"; import {WorkspaceFolderManager} from "../vscode-objs/WorkspaceFolderManager"; +// eslint-disable-next-line @typescript-eslint/naming-convention +const {NamedLogger} = logging; + +/** + * A compute target the user picked in the compute QuickPick. Structurally the + * `setup-local` compute shape, so python-setup can consume it directly without + * this module depending on the python-setup layer. Serverless is only ever + * returned version-complete (see {@link ConnectionCommands.selectServerless}). + */ +export type SelectedCompute = + | {kind: "cluster"; clusterId: string} + | {kind: "serverless"; version: string}; + function formatQuickPickClusterSize(sizeInMB: number): string { if (sizeInMB > 1024) { return Math.round(sizeInMB / 1024).toString() + " GB"; @@ -144,14 +157,17 @@ export class ConnectionCommands implements Disposable { } /** - * The returned handler resolves only once the picker has fully closed -- - * after any attach/enable it triggers has completed, or immediately when the - * user dismisses it. Callers that just open the picker can ignore the result - * (they always have); the resolution lets a caller that needs the outcome - * (python-setup, which re-reads the attached compute afterward) await it. + * The returned handler resolves once the picker has fully closed, to the + * compute the user attached -- a cluster or a version-complete serverless + * target -- or `undefined` when they dismissed it or picked "Create New + * Cluster" (which only opens the browser). Callers that just open the picker + * can ignore the result (they always have); returning the selection lets a + * caller that needs the outcome (python-setup) use it directly instead of + * re-reading the connection manager, whose cluster attach propagates + * asynchronously and would race an immediate read. */ attachClusterQuickPickCommand() { - return async (title?: string): Promise => { + return async (title?: string): Promise => { const workspaceClient = this.connectionManager.workspaceClient; const me = this.connectionManager.databricksWorkspace?.userName; if (!workspaceClient || !me) { @@ -213,28 +229,38 @@ export class ConnectionCommands implements Disposable { // Resolve only once the picker's work is done: on accept, after the // attach/enable (or browser hand-off) it triggers; on a plain hide, - // right away. `accepted` keeps the hide handler -- which also fires - // when accept disposes the picker -- from resolving early, before the - // accept branch's awaits have settled. - await new Promise((resolve) => { - let accepted = false; + // right away. `settled` guards both a repeated Enter (which would + // otherwise run concurrent attaches) and the hide handler -- which + // also fires when accept disposes the picker -- from resolving early, + // before the accept branch's awaits have settled. + return await new Promise((resolve) => { + let settled = false; quickPick.onDidAccept(async () => { - accepted = true; + if (settled) { + return; + } + settled = true; + let selected: SelectedCompute | undefined; try { const selectedItem = quickPick.selectedItems[0]; - if ("cluster" in selectedItem) { + if (selectedItem === undefined) { + // Accepted with nothing highlighted -- nothing to do. + } else if ("cluster" in selectedItem) { const cluster = selectedItem.cluster; 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()); - await this.selectServerless(); - return; + const version = await this.selectServerless(); + if (version !== undefined) { + selected = {kind: "serverless", version}; + } } else { await UrlUtils.openExternal( `${ @@ -245,17 +271,27 @@ export class ConnectionCommands implements Disposable { }#create/cluster` ); } - disposables.forEach((d) => d.dispose()); + } catch (e) { + // The attach/enable helpers surface their own failures + // (they are @onError-decorated); swallow here only so a + // throw can't become an unhandled rejection or leave the + // picker open. `selected` stays undefined. + NamedLogger.getOrCreate("Extension").error( + "Compute picker selection failed", + e + ); } finally { - resolve(); + disposables.forEach((d) => d.dispose()); + resolve(selected); } }); quickPick.onDidHide(() => { disposables.forEach((d) => d.dispose()); quickPick.dispose(); - if (!accepted) { - resolve(); + if (!settled) { + settled = true; + resolve(undefined); } }); }); @@ -269,18 +305,24 @@ 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 when serverless was enabled with one, or + * `undefined` -- when the version picker was dismissed (no change made), or + * when the feature is off (serverless is enabled, but version-less). Only + * the feature-on, version-complete case is a compute the caller can set up. */ - private async selectServerless() { + private async selectServerless(): Promise { 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; } /** diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index e69690f6d..f4047233a 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -58,6 +58,7 @@ import { resolveComputeFrom, } from "./python-setup/controllers/pythonSetupDeps"; import {PythonSetupDriftManager} from "./python-setup/controllers/PythonSetupDriftManager"; +import {SetupCompute} from "./python-setup/controllers/PythonSetupEnvironmentSetup"; import {resolveCliPath} from "./python-setup/utils/setupLocalArgs"; import { isPythonSetupEnabled, @@ -1005,14 +1006,14 @@ export async function activate( // setup flow. persistServerlessVersion: (version) => connectionManager.enableServerless(version), - // Reuses the existing compute picker. When the feature is opted in, - // its serverless branch also records the environment version, so a - // serverless selection made here comes back version-complete and - // setup need not re-prompt. Resolves once the picker closes; the - // setup flow then re-reads the attached compute. + // Reuses the existing compute picker, which returns the compute the + // user chose (or undefined if dismissed). When the feature is opted + // in, its serverless branch also resolves the environment version, + // so a serverless selection comes back version-complete and setup + // need not re-prompt. promptSelectCompute: () => Promise.resolve( - commands.executeCommand( + commands.executeCommand( "databricks.connection.attachClusterQuickPick" ) ), diff --git a/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.test.ts b/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.test.ts index 93a6134d3..0a641f154 100644 --- a/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.test.ts +++ b/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.test.ts @@ -7,6 +7,7 @@ import { resolveComputeFrom, } from "./pythonSetupDeps"; import {PythonSetupState} from "../../vscode-objs/StateStorage"; +import {SetupCompute} from "./PythonSetupEnvironmentSetup"; import {Telemetry} from "../../telemetry"; import { SUCCESS_DEFAULT, @@ -175,7 +176,7 @@ function makeWiring( }), promptServerlessVersion: async () => "4", persistServerlessVersion: async () => {}, - promptSelectCompute: async () => {}, + promptSelectCompute: async () => undefined, setActiveInterpreter: async () => {}, persistSetupState: () => {}, log: {append: () => {}, show: () => {}}, @@ -266,6 +267,7 @@ describe("makePythonSetupDeps resolveCompute", () => { promptSelectCompute: async () => { pickerOpened++; // User dismisses the picker without attaching anything. + return undefined; }, promptServerlessVersion: async () => { versionPrompted++; @@ -281,24 +283,21 @@ describe("makePythonSetupDeps resolveCompute", () => { expect(versionPrompted).to.equal(0); }); - it("runs against the compute the user attaches through the picker", async () => { - // The picker persists the selection through the connection manager, so - // resolveCompute re-reads the attachment after it closes. - let attached = { - serverless: false, - cluster: undefined as {id: string} | undefined, - serverlessVersion: undefined as string | undefined, - }; + it("runs against the cluster the picker returns, without re-reading state", async () => { + // The picker returns the chosen compute directly. Re-reading it from + // the connection manager would race the async, network-gated cluster + // attach -- so `attachedCompute` deliberately stays `none` here and the + // flow must still resolve to the returned cluster. const deps = makePythonSetupDeps( makeWiring({ - attachedCompute: () => attached, - promptSelectCompute: async () => { - attached = { - serverless: false, - cluster: {id: "c1"}, - serverlessVersion: undefined, - }; - }, + attachedCompute: () => ({ + serverless: false, + cluster: undefined, + serverlessVersion: undefined, + }), + promptSelectCompute: async (): Promise< + SetupCompute | undefined + > => ({kind: "cluster", clusterId: "c1"}), }) ); @@ -308,55 +307,55 @@ describe("makePythonSetupDeps resolveCompute", () => { }); }); - it("does not open the compute picker when a cluster is already attached", async () => { - let pickerOpened = 0; + it("runs against the serverless compute the picker returns", async () => { + // The picker's serverless branch resolves a version before enabling, so + // it comes back version-complete -- no follow-up version prompt. + let versionPrompted = 0; const deps = makePythonSetupDeps( makeWiring({ attachedCompute: () => ({ serverless: false, - cluster: {id: "c1"}, + cluster: undefined, serverlessVersion: undefined, }), - promptSelectCompute: async () => { - pickerOpened++; + promptSelectCompute: async (): Promise< + SetupCompute | undefined + > => ({kind: "serverless", version: "5"}), + promptServerlessVersion: async () => { + versionPrompted++; + return "4"; }, }) ); expect(await deps.resolveCompute()).to.deep.equal({ status: "ok", - compute: {kind: "cluster", clusterId: "c1"}, + compute: {kind: "serverless", version: "5"}, }); - expect(pickerOpened).to.equal(0); + expect(versionPrompted).to.equal(0); }); - it("prompts for a version when the picker attaches version-less serverless", async () => { - // The picker can leave serverless selected without a version (a config - // that enables serverless without opening the version sub-picker). The - // flow then resolves the version rather than dead-ending. - let attached = { - serverless: false, - cluster: undefined as {id: string} | undefined, - serverlessVersion: undefined as string | undefined, - }; + it("does not open the compute picker when a cluster is already attached", async () => { + let pickerOpened = 0; const deps = makePythonSetupDeps( makeWiring({ - attachedCompute: () => attached, + attachedCompute: () => ({ + serverless: false, + cluster: {id: "c1"}, + serverlessVersion: undefined, + }), promptSelectCompute: async () => { - attached = { - serverless: true, - cluster: undefined, - serverlessVersion: undefined, - }; + pickerOpened++; + return undefined; }, - promptServerlessVersion: async () => "4", }) ); expect(await deps.resolveCompute()).to.deep.equal({ status: "ok", - compute: {kind: "serverless", version: "4"}, + compute: {kind: "cluster", clusterId: "c1"}, }); + expect(pickerOpened).to.equal(0); }); it("never prompts when serverless already has a version", async () => { diff --git a/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.ts b/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.ts index 006df2aa9..e2fa354cb 100644 --- a/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.ts +++ b/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.ts @@ -139,12 +139,14 @@ export interface PythonSetupWiringDeps { persistServerlessVersion: (version: string) => Promise; /** * Open the compute picker so the user can attach a target when none is - * selected, resolving once it closes. The selection is persisted through - * the connection manager (a cluster attaches, serverless enables), so the - * caller re-reads {@link attachedCompute} afterward rather than taking a - * return value -- and re-reads `none` when the user dismissed it. + * selected, resolving to the chosen compute once it closes (or `undefined` + * when the user dismisses it). The picker also attaches the selection + * through the connection manager as a side effect, but the caller uses this + * return value directly rather than re-reading {@link attachedCompute}: a + * cluster attach propagates asynchronously (a network-gated, + * fire-and-forget config listener), so an immediate re-read would race it. */ - promptSelectCompute: () => Promise; + promptSelectCompute: () => Promise; /** Point the MS Python extension at an interpreter path (project-scoped). */ setActiveInterpreter: (interpreterPath: string, root: Uri) => Promise; /** Persist the post-setup state (workspace-scoped) for drift detection. */ @@ -177,16 +179,19 @@ export function makePythonSetupDeps( projectRoot: wiring.projectRoot, isVisible, resolveCompute: async () => { - let resolution = resolveComputeFrom(wiring.attachedCompute()); + const resolution = resolveComputeFrom(wiring.attachedCompute()); if (resolution.status === "none") { // Nothing attached: open the compute picker inline so the user - // can choose a target, rather than dead-ending the CTA. The - // picker persists the selection through the connection manager, - // so re-read the attachment afterward -- it reflects whatever - // they chose, or stays `none` if they dismissed the picker (in - // which case the orchestrator shows the guidance message). - await wiring.promptSelectCompute(); - resolution = resolveComputeFrom(wiring.attachedCompute()); + // can choose a target, rather than dead-ending the CTA. Use the + // picker's own return value -- re-reading the attachment from + // the connection manager would race the cluster attach, which + // propagates asynchronously through a network-gated config + // listener. `undefined` means the user dismissed the picker, so + // fall through to `none` and let the orchestrator guide them. + const picked = await wiring.promptSelectCompute(); + return picked === undefined + ? {status: "none"} + : {status: "ok", compute: picked}; } if (resolution.status !== "needsServerlessVersion") { return resolution; From 53be6610af0f8d755e781569fa726a583b3e9814 Mon Sep 17 00:00:00 2001 From: "@rugpanov" Date: Thu, 13 Aug 2026 23:51:22 +0200 Subject: [PATCH 3/5] test(python-setup): cover the picker's return contract; tie its compute type Why: Round-1 multi-source review (Codex + a Claude reviewer) surfaced three points. Codex's "attach failure returns as success" was verified a false positive: setup-local receives the clusterId/version descriptor directly (setupLocalArgs), so it is valid regardless of whether the config write persisted, and @onError already surfaces any failure to the user. The two actionable notes were a test-coverage gap on the picker's producer logic and an unchecked type link. What: - Add unit tests for attachClusterQuickPickCommand's return contract via a scriptable QuickPick fake: resolves to the attached cluster (attaching once), undefined on dismissal, undefined on empty accept, and the `settled` re-entry guard (a second Enter neither re-attaches nor changes the result). This is the riskiest new code and had no producer-side test. - Alias SelectedCompute to SetupCompute (the setup-local compute shape) instead of re-declaring it, so drift is a compile error rather than a silent mismatch through executeCommand's unchecked generic. - Use Loggers.Extension (the enum) instead of the "Extension" string literal, matching the codebase convention. Verification: - yarn build; full unit suite green (740 passing, +4 new). yarn fix / test:lint clean. Co-authored-by: Isaac --- .../configuration/ConnectionCommands.test.ts | 139 ++++++++++++++++++ .../src/configuration/ConnectionCommands.ts | 18 ++- 2 files changed, 149 insertions(+), 8 deletions(-) diff --git a/packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts b/packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts index 626557093..f4190094d 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts +++ b/packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts @@ -3,11 +3,59 @@ 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"; +/** + * A scriptable stand-in for the compute QuickPick. Unlike a fire-on-show fake, + * the test drives {@link accept}/{@link hide} explicitly -- because + * `attachClusterQuickPickCommand` registers its `onDidAccept`/`onDidHide` + * handlers only after calling `show()`, so firing them at show time would race + * (and never resolve) the command's Promise. + */ +class FakeComputeQuickPick { + title?: string; + keepScrollPosition = false; + busy = false; + canSelectMany = false; + items: readonly QuickPickItem[] = []; + selectedItems: readonly QuickPickItem[] = []; + disposed = false; + private readonly acceptCbs: Array<() => void | Promise> = []; + private readonly hideCbs: Array<() => void> = []; + + onDidAccept(cb: () => void | Promise) { + 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 +82,95 @@ 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 () => { + // The `settled` flag must short-circuit a repeated accept -- otherwise + // a double Enter would run concurrent attaches and could resolve 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, []); + }); + }); }); diff --git a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts b/packages/databricks-vscode/src/configuration/ConnectionCommands.ts index 1689d0fbd..6ced80d7c 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts +++ b/packages/databricks-vscode/src/configuration/ConnectionCommands.ts @@ -30,20 +30,22 @@ 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 the user picked in the compute QuickPick. Structurally the - * `setup-local` compute shape, so python-setup can consume it directly without - * this module depending on the python-setup layer. Serverless is only ever - * returned version-complete (see {@link ConnectionCommands.selectServerless}). + * A compute target the user picked in the compute QuickPick. Aliased to the + * `setup-local` compute shape (rather than re-declared) so python-setup can + * consume the picker's result directly and any drift in that shape is a compile + * error here, not a silent runtime mismatch through `executeCommand`'s unchecked + * generic. Serverless is only ever returned version-complete (see + * {@link ConnectionCommands.selectServerless}). */ -export type SelectedCompute = - | {kind: "cluster"; clusterId: string} - | {kind: "serverless"; version: string}; +export type SelectedCompute = SetupCompute; function formatQuickPickClusterSize(sizeInMB: number): string { if (sizeInMB > 1024) { @@ -276,7 +278,7 @@ export class ConnectionCommands implements Disposable { // (they are @onError-decorated); swallow here only so a // throw can't become an unhandled rejection or leave the // picker open. `selected` stays undefined. - NamedLogger.getOrCreate("Extension").error( + NamedLogger.getOrCreate(Loggers.Extension).error( "Compute picker selection failed", e ); From 8b244c2b1bfb04e830729444e8f6f0566238407d Mon Sep 17 00:00:00 2001 From: "@rugpanov" Date: Fri, 14 Aug 2026 07:48:08 +0200 Subject: [PATCH 4/5] docs(python-setup): document the picker's best-effort attach contract + tests Why: Iterative review re-raised that the picker returns the chosen compute even when attachCluster/enableServerless silently fail (they are @onError with throw:false, so a config-write failure shows a popup but does not throw). Decision: keep this best-effort -- the returned descriptor is the compute the user picked and is all setup-local needs (it takes the id/version directly), and the failure is already surfaced to the user. The prior code comment overclaimed that the catch nulls the selection on attach failure, which is dead for those swallowed operations. What: - Correct the catch comment: it is defense-in-depth for an UNEXPECTED throw (version sub-picker, openExternal, a future non-decorated path), not a success/failure discriminator for the @onError-swallowed attach/enable. - Document at the cluster branch why returning the descriptor is intentional even when the attach's persistence fails. - Add contract tests: the chosen cluster is still returned when the attach silently fails (best-effort), and an unexpected throw settles the Promise to undefined (no hang, no selection). Verification: - yarn build; full unit suite green (742 passing, +2). yarn fix / test:lint clean. Co-authored-by: Isaac --- .../configuration/ConnectionCommands.test.ts | 70 +++++++++++++++++++ .../src/configuration/ConnectionCommands.ts | 19 +++-- 2 files changed, 85 insertions(+), 4 deletions(-) diff --git a/packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts b/packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts index f4190094d..5e5054532 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts +++ b/packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts @@ -172,5 +172,75 @@ describe(__filename, () => { 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 config write + // surfaces its own popup and resolves without throwing, so from the + // picker's side it looks exactly like success. By design we still + // return the picked target -- it is what the user chose and all + // setup-local needs -- rather than dropping the selection. + const attach = async () => { + // Resolves (does not throw) even though the "attach" failed, + // mimicking the @onError-swallowed path. + }; + 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 () => { + // Defense-in-depth: a throw from any step must settle the Promise + // (not hang) and yield no selection. + 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); + }); }); }); diff --git a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts b/packages/databricks-vscode/src/configuration/ConnectionCommands.ts index 6ced80d7c..7483e08b2 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts +++ b/packages/databricks-vscode/src/configuration/ConnectionCommands.ts @@ -249,6 +249,13 @@ export class ConnectionCommands implements Disposable { // Accepted with nothing highlighted -- nothing to do. } else if ("cluster" in selectedItem) { const cluster = selectedItem.cluster; + // Best-effort attach: attachCluster is @onError with + // throw:false, so a failed config write surfaces its + // own popup and does NOT throw here. We still return + // the chosen target -- it is the compute the user + // picked and is all setup-local needs (it takes the + // id directly); the attach's persistence is a + // separate concern already surfaced to the user. await this.connectionManager.attachCluster( cluster.id ); @@ -274,10 +281,14 @@ export class ConnectionCommands implements Disposable { ); } } catch (e) { - // The attach/enable helpers surface their own failures - // (they are @onError-decorated); swallow here only so a - // throw can't become an unhandled rejection or leave the - // picker open. `selected` stays undefined. + // Defense-in-depth. The attach/enable helpers are + // @onError(throw:false) and surface their own failures + // without throwing, so they do not reach here; this + // guards an UNEXPECTED throw from another step (the + // serverless version sub-picker, openExternal, or a + // future non-decorated path) so it can't become an + // unhandled rejection or leave the Promise pending. + // `selected` stays undefined in that case. NamedLogger.getOrCreate(Loggers.Extension).error( "Compute picker selection failed", e From 001396e468081b98d6a5ecbffa72b2340f917ea1 Mon Sep 17 00:00:00 2001 From: "@rugpanov" Date: Fri, 14 Aug 2026 07:57:44 +0200 Subject: [PATCH 5/5] style(python-setup): trim the added comments to be concise Why: The explanatory comments added across the feature were far longer than the code warranted. What: - Condense the doc/inline comments in ConnectionCommands.ts (picker command, settle guard, best-effort attach, catch, SelectedCompute, selectServerless), pythonSetupDeps.ts (promptSelectCompute, resolveCompute), extension.ts, and the two test files. No code or behavior changes. Verification: - yarn build; full unit suite green (742 passing). yarn fix / test:lint clean. Co-authored-by: Isaac --- .../configuration/ConnectionCommands.test.ts | 26 +++----- .../src/configuration/ConnectionCommands.ts | 60 +++++++------------ packages/databricks-vscode/src/extension.ts | 7 +-- .../controllers/pythonSetupDeps.test.ts | 19 ++---- .../controllers/pythonSetupDeps.ts | 21 +++---- 5 files changed, 43 insertions(+), 90 deletions(-) diff --git a/packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts b/packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts index 5e5054532..0ca8d54d3 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts +++ b/packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts @@ -11,11 +11,9 @@ import { } from "./ConnectionCommands"; /** - * A scriptable stand-in for the compute QuickPick. Unlike a fire-on-show fake, - * the test drives {@link accept}/{@link hide} explicitly -- because - * `attachClusterQuickPickCommand` registers its `onDidAccept`/`onDidHide` - * handlers only after calling `show()`, so firing them at show time would race - * (and never resolve) the command's Promise. + * 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; @@ -152,8 +150,7 @@ describe(__filename, () => { }); it("guards re-entry: a second Enter neither re-attaches nor changes the result", async () => { - // The `settled` flag must short-circuit a repeated accept -- otherwise - // a double Enter would run concurrent attaches and could resolve twice. + // Without the `settled` guard a double Enter would attach twice. const resultP = commands.attachClusterQuickPickCommand()(); await fakePick.accept([clusterItem("c1")]); await fakePick.accept([clusterItem("c2")]); @@ -174,15 +171,9 @@ describe(__filename, () => { }); it("still returns the chosen cluster when the attach silently fails (best-effort)", async () => { - // attachCluster is @onError(throw:false): a failed config write - // surfaces its own popup and resolves without throwing, so from the - // picker's side it looks exactly like success. By design we still - // return the picked target -- it is what the user chose and all - // setup-local needs -- rather than dropping the selection. - const attach = async () => { - // Resolves (does not throw) even though the "attach" failed, - // mimicking the @onError-swallowed path. - }; + // 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, { @@ -213,8 +204,7 @@ describe(__filename, () => { }); it("resolves to undefined if an unexpected error is thrown during selection", async () => { - // Defense-in-depth: a throw from any step must settle the Promise - // (not hang) and yield no selection. + // A throw from any step must settle the Promise, not hang. const cmds = new ConnectionCommands( {} as never, { diff --git a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts b/packages/databricks-vscode/src/configuration/ConnectionCommands.ts index 7483e08b2..2ee19cfe3 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts +++ b/packages/databricks-vscode/src/configuration/ConnectionCommands.ts @@ -38,12 +38,9 @@ import {Loggers} from "../logger"; const {NamedLogger} = logging; /** - * A compute target the user picked in the compute QuickPick. Aliased to the - * `setup-local` compute shape (rather than re-declared) so python-setup can - * consume the picker's result directly and any drift in that shape is a compile - * error here, not a silent runtime mismatch through `executeCommand`'s unchecked - * generic. Serverless is only ever returned version-complete (see - * {@link ConnectionCommands.selectServerless}). + * 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; @@ -159,14 +156,10 @@ export class ConnectionCommands implements Disposable { } /** - * The returned handler resolves once the picker has fully closed, to the - * compute the user attached -- a cluster or a version-complete serverless - * target -- or `undefined` when they dismissed it or picked "Create New - * Cluster" (which only opens the browser). Callers that just open the picker - * can ignore the result (they always have); returning the selection lets a - * caller that needs the outcome (python-setup) use it directly instead of - * re-reading the connection manager, whose cluster attach propagates - * asynchronously and would race an immediate read. + * 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): Promise => { @@ -229,12 +222,9 @@ export class ConnectionCommands implements Disposable { refreshQuickPickItems(); quickPick.show(); - // Resolve only once the picker's work is done: on accept, after the - // attach/enable (or browser hand-off) it triggers; on a plain hide, - // right away. `settled` guards both a repeated Enter (which would - // otherwise run concurrent attaches) and the hide handler -- which - // also fires when accept disposes the picker -- from resolving early, - // before the accept branch's awaits have settled. + // `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((resolve) => { let settled = false; quickPick.onDidAccept(async () => { @@ -246,16 +236,13 @@ export class ConnectionCommands implements Disposable { try { const selectedItem = quickPick.selectedItems[0]; if (selectedItem === undefined) { - // Accepted with nothing highlighted -- nothing to do. + // Accepted with nothing highlighted. } else if ("cluster" in selectedItem) { const cluster = selectedItem.cluster; - // Best-effort attach: attachCluster is @onError with - // throw:false, so a failed config write surfaces its - // own popup and does NOT throw here. We still return - // the chosen target -- it is the compute the user - // picked and is all setup-local needs (it takes the - // id directly); the attach's persistence is a - // separate concern already surfaced to the user. + // 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 ); @@ -281,14 +268,9 @@ export class ConnectionCommands implements Disposable { ); } } catch (e) { - // Defense-in-depth. The attach/enable helpers are - // @onError(throw:false) and surface their own failures - // without throwing, so they do not reach here; this - // guards an UNEXPECTED throw from another step (the - // serverless version sub-picker, openExternal, or a - // future non-decorated path) so it can't become an - // unhandled rejection or leave the Promise pending. - // `selected` stays undefined in that case. + // 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 @@ -319,10 +301,8 @@ export class ConnectionCommands implements Disposable { * picker, no compute change is made. With the feature off this is the * plain, unchanged serverless enable. * - * Returns the confirmed version when serverless was enabled with one, or - * `undefined` -- when the version picker was dismissed (no change made), or - * when the feature is off (serverless is enabled, but version-less). Only - * the feature-on, version-complete case is a compute the caller can set up. + * Returns the confirmed version, or `undefined` if the picker was dismissed + * or the feature is off (serverless enabled but version-less). */ private async selectServerless(): Promise { if (!isPythonSetupEnabled()) { diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index f4047233a..906d0161e 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -1006,11 +1006,8 @@ export async function activate( // setup flow. persistServerlessVersion: (version) => connectionManager.enableServerless(version), - // Reuses the existing compute picker, which returns the compute the - // user chose (or undefined if dismissed). When the feature is opted - // in, its serverless branch also resolves the environment version, - // so a serverless selection comes back version-complete and setup - // need not re-prompt. + // Reuses the compute picker, which returns the chosen compute (or + // undefined if dismissed); its serverless branch is version-complete. promptSelectCompute: () => Promise.resolve( commands.executeCommand( diff --git a/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.test.ts b/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.test.ts index 0a641f154..2386b0736 100644 --- a/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.test.ts +++ b/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.test.ts @@ -257,17 +257,14 @@ describe("makePythonSetupDeps resolveCompute", () => { }); it("offers the compute picker (not the version picker) when nothing is attached", async () => { - // Nothing selected is a missing target, not a missing detail: the flow - // opens the compute picker so the user can choose one -- it must never - // jump straight to the serverless version picker. + // Nothing attached opens the compute picker, never the version picker. let versionPrompted = 0; let pickerOpened = 0; const deps = makePythonSetupDeps( makeWiring({ promptSelectCompute: async () => { pickerOpened++; - // User dismisses the picker without attaching anything. - return undefined; + return undefined; // dismissed }, promptServerlessVersion: async () => { versionPrompted++; @@ -276,18 +273,15 @@ describe("makePythonSetupDeps resolveCompute", () => { }) ); - // Dismissed with nothing attached -> still the dead-end `none`, which - // the orchestrator turns into the guidance message. + // Dismissed -> `none`, which the orchestrator turns into guidance. expect(await deps.resolveCompute()).to.deep.equal({status: "none"}); expect(pickerOpened).to.equal(1); expect(versionPrompted).to.equal(0); }); it("runs against the cluster the picker returns, without re-reading state", async () => { - // The picker returns the chosen compute directly. Re-reading it from - // the connection manager would race the async, network-gated cluster - // attach -- so `attachedCompute` deliberately stays `none` here and the - // flow must still resolve to the returned cluster. + // Uses the picker's return value; `attachedCompute` stays `none` to + // prove the flow doesn't re-read (which would race the attach). const deps = makePythonSetupDeps( makeWiring({ attachedCompute: () => ({ @@ -308,8 +302,7 @@ describe("makePythonSetupDeps resolveCompute", () => { }); it("runs against the serverless compute the picker returns", async () => { - // The picker's serverless branch resolves a version before enabling, so - // it comes back version-complete -- no follow-up version prompt. + // The picker returns serverless version-complete -- no follow-up prompt. let versionPrompted = 0; const deps = makePythonSetupDeps( makeWiring({ diff --git a/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.ts b/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.ts index e2fa354cb..7de8b7b94 100644 --- a/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.ts +++ b/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.ts @@ -138,13 +138,10 @@ export interface PythonSetupWiringDeps { */ persistServerlessVersion: (version: string) => Promise; /** - * Open the compute picker so the user can attach a target when none is - * selected, resolving to the chosen compute once it closes (or `undefined` - * when the user dismisses it). The picker also attaches the selection - * through the connection manager as a side effect, but the caller uses this - * return value directly rather than re-reading {@link attachedCompute}: a - * cluster attach propagates asynchronously (a network-gated, - * fire-and-forget config listener), so an immediate re-read would race it. + * Open the compute picker when nothing is attached, resolving to the chosen + * compute (or `undefined` if dismissed). The caller uses this return value + * directly, not a re-read of {@link attachedCompute}: a cluster attach + * propagates asynchronously, so an immediate re-read would race it. */ promptSelectCompute: () => Promise; /** Point the MS Python extension at an interpreter path (project-scoped). */ @@ -181,13 +178,9 @@ export function makePythonSetupDeps( resolveCompute: async () => { const resolution = resolveComputeFrom(wiring.attachedCompute()); if (resolution.status === "none") { - // Nothing attached: open the compute picker inline so the user - // can choose a target, rather than dead-ending the CTA. Use the - // picker's own return value -- re-reading the attachment from - // the connection manager would race the cluster attach, which - // propagates asynchronously through a network-gated config - // listener. `undefined` means the user dismissed the picker, so - // fall through to `none` and let the orchestrator guide them. + // Nothing attached: offer the picker inline instead of + // dead-ending. Use its return value (not a racy re-read); + // `undefined` means dismissed, so fall through to `none`. const picked = await wiring.promptSelectCompute(); return picked === undefined ? {status: "none"}