[8/?] feat(python-setup): merge serverless-version selection into the compute picker

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

@@ -23,6 +23,15 @@ import {
 } from "../ui/configuration-view/AuthTypeComponent";
 import {ManualLoginSource} from "../telemetry/constants";
 import {onError} from "../utils/onErrorDecorator";
+import {
+    isPythonSetupEnabled,
+    resolveServerlessVersion,
+} from "../python-setup/utils/serverlessVersionResolver";
+import {pickServerlessVersion} from "../python-setup/utils/serverlessVersionPicker";
+import {collectBundleServerlessVersions} from "../python-setup/utils/bundleServerlessVersions";
+import {collectProjectNotebookVersions} from "../python-setup/utils/projectNotebookVersions";
+import {VersionObservation} from "../python-setup/utils/serverlessVersionScoring";
+import {WorkspaceFolderManager} from "../vscode-objs/WorkspaceFolderManager";
 
 function formatQuickPickClusterSize(sizeInMB: number): string {
     if (sizeInMB > 1024) {
@@ -63,7 +72,8 @@ export class ConnectionCommands implements Disposable {
         private connectionManager: ConnectionManager,
         private readonly clusterModel: ClusterModel,
         private readonly configModel: ConfigModel,
-        private readonly cli: CliWrapper
+        private readonly cli: CliWrapper,
+        private readonly workspaceFolderManager: WorkspaceFolderManager
     ) {}
 
     /**
@@ -187,7 +197,11 @@ export class ConnectionCommands implements Disposable {
                     const cluster = selectedItem.cluster;
                     await this.connectionManager.attachCluster(cluster.id);
                 } else if (selectedItem.label === "$(cloud) Serverless") {
-                    await this.connectionManager.enableServerless();
+                    // 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(
                         `${
@@ -208,6 +222,77 @@ export class ConnectionCommands implements Disposable {
         };
     }
 
+    /**
+     * Enable serverless compute. When the uv-native python-setup feature is
+     * opted into, first ask the user to confirm the serverless environment
+     * version (ranked from the project's bundle) and persist it with the
+     * 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.
+     */
+    private async selectServerless() {
+        if (!isPythonSetupEnabled()) {
+            await this.connectionManager.enableServerless();
+            return;
+        }
+        const version = await this.pickServerlessVersion();
+        if (version === undefined) {
+            // User dismissed the version picker -- don't switch compute.
+            return;
+        }
+        await this.connectionManager.enableServerless(version);
+    }
+
+    /**
+     * Resolve the serverless environment version: gather the project's version
+     * evidence, score it, and let the user confirm the best-ranked candidate.
+     * Returns the confirmed bare version, or undefined if dismissed. Delegates
+     * to {@link resolveServerlessVersion} so the collect->score->pick pipeline
+     * lives in one place; this call site only supplies where the evidence comes
+     * from (the bundle today) and how the user confirms it.
+     */
+    private async pickServerlessVersion(): Promise<string | undefined> {
+        return resolveServerlessVersion({
+            collectObservations: () => this.collectServerlessObservations(),
+            pick: pickServerlessVersion,
+        });
+    }
+
+    /**
+     * Gather serverless-version evidence from the project's local sources
+     * (bundle config + notebooks). Each source is collected independently and
+     * guarded, so one failing source never blocks the other or compute
+     * selection; the scorer merges and de-dupes across sources. (The
+     * workspace-default API source is not yet available — see DECO-27782.)
+     */
+    private async collectServerlessObservations(): Promise<
+        VersionObservation[]
+    > {
+        const [bundle, notebooks] = await Promise.all([
+            (async () => {
+                try {
+                    const validateConfig =
+                        await this.configModel.get("validateConfig");
+                    return collectBundleServerlessVersions(validateConfig);
+                } catch {
+                    return [] as VersionObservation[];
+                }
+            })(),
+            (async () => {
+                try {
+                    const projectRoot =
+                        this.workspaceFolderManager.activeProjectUri.fsPath;
+                    return await collectProjectNotebookVersions(projectRoot);
+                } catch {
+                    // No active project, or notebook scan failed -- contribute
+                    // nothing rather than blocking compute selection.
+                    return [] as VersionObservation[];
+                }
+            })(),
+        ]);
+        return [...bundle, ...notebooks];
+    }
+
     /**
      * Set cluster to undefined and remove cluster ID from settings file
      */

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

@@ -26,6 +26,7 @@ import {AutoLoginSource, ManualLoginSource} from "../telemetry/constants";
 import {Barrier} from "../locking/Barrier";
 import {WorkspaceFolderManager} from "../vscode-objs/WorkspaceFolderManager";
 import {ProjectConfigFile} from "../file-managers/ProjectConfigFile";
+import {isSupportedVersion} from "../python-setup/utils/serverlessVersionScoring";
 
 // eslint-disable-next-line @typescript-eslint/naming-convention
 const {NamedLogger} = logging;
@@ -50,6 +51,7 @@ export class ConnectionManager implements Disposable {
     private _databricksWorkspace?: DatabricksWorkspace;
     private _metadataService: MetadataService;
     private _serverlessEnabled: boolean = false;
+    private _serverlessVersion: string | undefined;
 
     private readonly onDidChangeStateEmitter: EventEmitter<ConnectionState> =
         new EventEmitter();
@@ -158,7 +160,21 @@ export class ConnectionManager implements Disposable {
         const autoEnable =
             serverless === undefined && !this.cluster && computeId === "auto";
         if (serverless || autoEnable) {
-            await this.enableServerless();
+            // Load the persisted version (if any) so it survives a reload
+            // without re-prompting. A version-less serverless config (older
+            // extensions, or serverless enabled before a version was chosen)
+            // leaves this undefined, and consumers fall back to their default.
+            // The stored value is untrusted (config files are hand-editable and
+            // may predate the supported range), so re-validate here rather than
+            // exposing a value the `--serverless-version` flag would reject --
+            // an unsupported version is dropped to undefined (scored default).
+            const storedVersion =
+                await this.configModel.get("serverlessVersion");
+            this._serverlessVersion =
+                storedVersion !== undefined && isSupportedVersion(storedVersion)
+                    ? storedVersion
+                    : undefined;
+            await this.enableServerless(this._serverlessVersion);
         } else {
             await this.disableServerless();
         }
@@ -228,6 +244,16 @@ export class ConnectionManager implements Disposable {
         return this._serverlessEnabled;
     }
 
+    /**
+     * The persisted serverless environment version (the CLI's bare-integer
+     * `--serverless-version` value, e.g. "5"), or undefined when serverless is
+     * not selected or no version has been chosen yet. Consumers treat undefined
+     * as "fall back to the scored default".
+     */
+    get serverlessVersion(): string | undefined {
+        return this._serverlessEnabled ? this._serverlessVersion : undefined;
+    }
+
     get syncDestinationMapper(): SyncDestinationMapper | undefined {
         return this._syncDestinationMapper;
     }
@@ -465,7 +491,15 @@ export class ConnectionManager implements Disposable {
     @onError({
         popup: {prefix: "Failed to enable serverless mode."},
     })
-    async enableServerless() {
+    async enableServerless(version?: string) {
+        // Persist the version whenever one is supplied, even if serverless is
+        // already enabled -- this is how re-picking the version (without
+        // toggling compute) is saved. `undefined` leaves any existing persisted
+        // version in place rather than clearing it.
+        if (version !== undefined && version !== this._serverlessVersion) {
+            this._serverlessVersion = version;
+            await this.configModel.set("serverlessVersion", version);
+        }
         if (!this._serverlessEnabled) {
             this._serverlessEnabled = true;
             await this.configModel.set("serverless", true);
@@ -485,7 +519,12 @@ export class ConnectionManager implements Disposable {
     async disableServerless() {
         if (this._serverlessEnabled) {
             this._serverlessEnabled = false;
+            // Clear the version too: it only has meaning while serverless is
+            // the selected compute, and leaving a stale value would let it
+            // resurface if serverless is re-enabled later without re-picking.
+            this._serverlessVersion = undefined;
             await this.configModel.set("serverless", false);
+            await this.configModel.set("serverlessVersion", undefined);
             this.customWhenContext.setServerless(false);
             this.onDidChangeClusterEmitter.fire(undefined);
         }

packages/databricks-vscode/src/configuration/models/OverrideableConfigModel.test.ts

@@ -0,0 +1,85 @@
+import {expect} from "chai";
+import * as tmp from "tmp";
+import path from "node:path";
+import {readFileSync, writeFileSync} from "node:fs";
+import {Uri} from "vscode";
+import {
+    isOverrideableConfigKey,
+    OverrideableConfigModel,
+} from "./OverrideableConfigModel";
+
+describe("OverrideableConfigModel serverlessVersion", () => {
+    const cleanups: Array<() => void> = [];
+    afterEach(() => {
+        while (cleanups.length) {
+            cleanups.pop()!();
+        }
+    });
+
+    function tempStorageFile(): Uri {
+        const dir = tmp.dirSync({unsafeCleanup: true});
+        cleanups.push(dir.removeCallback);
+        return Uri.file(path.join(dir.name, "vscode.overrides.json"));
+    }
+
+    it("treats serverlessVersion as an overrideable key", () => {
+        expect(isOverrideableConfigKey("serverlessVersion")).to.equal(true);
+    });
+
+    it("persists a serverless version alongside the serverless flag", async () => {
+        const file = tempStorageFile();
+
+        await OverrideableConfigModel._write(file, "serverless", "dev", true);
+        await OverrideableConfigModel._write(
+            file,
+            "serverlessVersion",
+            "dev",
+            "5"
+        );
+
+        const data = JSON.parse(readFileSync(file.fsPath, "utf-8"));
+        expect(data.serverless).to.equal(true);
+        expect(data.serverlessVersion).to.equal("5");
+    });
+
+    it("clears the version when written undefined (revert to fallback)", async () => {
+        const file = tempStorageFile();
+        await OverrideableConfigModel._write(
+            file,
+            "serverlessVersion",
+            "dev",
+            "4"
+        );
+
+        await OverrideableConfigModel._write(
+            file,
+            "serverlessVersion",
+            "dev",
+            undefined
+        );
+
+        const data = JSON.parse(readFileSync(file.fsPath, "utf-8"));
+        expect(data.serverlessVersion).to.equal(undefined);
+    });
+
+    it("leaves a legacy version-less serverless config untouched (backward compatible)", async () => {
+        const file = tempStorageFile();
+        // A config written by an older extension: serverless on, no version.
+        writeFileSync(
+            file.fsPath,
+            JSON.stringify({serverless: true, clusterId: "abc"})
+        );
+
+        // Writing an unrelated key must not fabricate a serverlessVersion.
+        await OverrideableConfigModel._write(
+            file,
+            "authProfile",
+            "dev",
+            "DEFAULT"
+        );
+
+        const data = JSON.parse(readFileSync(file.fsPath, "utf-8"));
+        expect(data.serverless).to.equal(true);
+        expect(data).to.not.have.property("serverlessVersion");
+    });
+});

packages/databricks-vscode/src/configuration/models/OverrideableConfigModel.ts

@@ -10,6 +10,17 @@ export type OverrideableConfigState = {
     authProfile?: string;
     clusterId?: string;
     serverless?: boolean;
+    /**
+     * The serverless environment version to provision the local Python
+     * environment against (the CLI's bare-integer `--serverless-version` value,
+     * e.g. "5"). Chosen once alongside the serverless compute selection and
+     * persisted so setup need not re-prompt.
+     *
+     * Optional and independent of `serverless`: existing configs that predate
+     * this field have `serverless: true` with no version, which stays valid --
+     * a version-less serverless selection falls back to the scored default.
+     */
+    serverlessVersion?: string;
     useClusterOverride?: boolean;
 };
 
@@ -21,6 +32,7 @@ export function isOverrideableConfigKey(
         "clusterId",
         "useClusterOverride",
         "serverless",
+        "serverlessVersion",
     ].includes(key);
 }
 

packages/databricks-vscode/src/extension.ts

@@ -787,7 +787,8 @@ export async function activate(
         connectionManager,
         clusterModel,
         configModel,
-        cli
+        cli,
+        workspaceFolderManager
     );
 
     context.subscriptions.push(