[Bug Bash] Frictionless VSCode Python Setup — default-on build

packages/databricks-vscode/package.json

@@ -1583,7 +1583,7 @@
                     },
                     "databricks.experiments.optInto": {
                         "type": "array",
-                        "default": [],
+                        "default": ["environment.pythonSetup"],
                         "items": {
                             "enum": [
                                 "views.cluster",

packages/databricks-vscode/src/extension.ts

@@ -16,6 +16,7 @@ import {ClusterListDataProvider} from "./cluster/ClusterListDataProvider";
 import {ClusterModel} from "./cluster/ClusterModel";
 import {ClusterCommands} from "./cluster/ClusterCommands";
 import {ConfigurationDataProvider} from "./ui/configuration-view/ConfigurationDataProvider";
+import {composePythonSetupEntry} from "./ui/configuration-view/pythonSetupEntry";
 import {AiToolsManager} from "./aitools/AiToolsManager";
 import {AiToolsCommands} from "./aitools/AiToolsCommands";
 import {RunCommands} from "./run/RunCommands";
@@ -51,7 +52,11 @@ import {
 import {PythonSetupManagerDetector} from "./python-setup/utils/PythonSetupManagerDetector";
 import {PythonSetupCliClient} from "./python-setup/gateways/PythonSetupCliClient";
 import {PythonSetupEnvironmentSetup} from "./python-setup/controllers/PythonSetupEnvironmentSetup";
-import {makePythonSetupDeps} from "./python-setup/controllers/pythonSetupDeps";
+import {
+    makePythonSetupDeps,
+    resolveComputeFrom,
+} from "./python-setup/controllers/pythonSetupDeps";
+import {PythonSetupDriftManager} from "./python-setup/controllers/PythonSetupDriftManager";
 import {resolveCliPath} from "./python-setup/utils/setupLocalArgs";
 import {
     isPythonSetupEnabled,
@@ -1014,15 +1019,112 @@ export async function activate(
             pythonSetupEnvironment.setup,
             pythonSetupEnvironment
         ),
-        // Re-run affordance on the "Python environment ready" row. Delegates to
-        // the same setup handler (re-entrancy-guarded); a distinct id lets the
-        // menu show a "Re-run Python setup" title instead of the initial one.
+        // Re-run affordance shared by the "Python environment ready" row and the
+        // drifted row. Delegates to the same setup handler (re-entrancy-guarded);
+        // a distinct command id lets the menu show a "Re-run Python setup" title
+        // and gives re-runs their own COMMAND_EXECUTION telemetry.
         telemetry.registerCommand(
             "databricks.environment.rerunPythonEnv",
             pythonSetupEnvironment.setup,
             pythonSetupEnvironment
         )
     );
+    // Drives the config-view row's "out of date" state: on compute/open/setup
+    // triggers it silently resolves the selected compute's env key via a CLI
+    // dry-run and compares it against the last successful setup's key. Every
+    // resolution path is fail-safe (returns undefined => "unknown", no drift) and
+    // never surfaces UI.
+    const pythonSetupDrift = new PythonSetupDriftManager({
+        // Reuse the exact feature+greenfield gate the row is shown under.
+        isVisible: () => pythonSetupEnvironment.isVisible(),
+        getPersistedEnvKey: () =>
+            stateStorage.get("databricks.pythonSetup.setupState")?.envKey,
+        // Cheap, synchronous compute identity (no CLI). A cluster's env key is
+        // derived from its Spark version, so include it: a runtime-state change
+        // (RUNNING -> TERMINATED) keeps the descriptor stable and is skipped,
+        // while a DBR edit changes it and re-checks. undefined means nothing
+        // comparable is attached (drift is then meaningless).
+        getComputeDescriptor: () => {
+            const cluster = connectionManager.cluster;
+            if (cluster) {
+                return `cluster:${cluster.id}:${cluster.sparkVersion}`;
+            }
+            if (connectionManager.serverless) {
+                const version = connectionManager.serverlessVersion;
+                return version === undefined
+                    ? undefined
+                    : `serverless:${version}`;
+            }
+            return undefined;
+        },
+        resolveCurrentEnvKey: async (token) => {
+            // activeProjectUri throws when no project is active; degrade to
+            // "unknown" rather than letting it reject into the drift check.
+            let root: string | undefined;
+            try {
+                root = workspaceFolderManager.activeProjectUri.fsPath;
+            } catch {
+                return undefined;
+            }
+            const resolution = resolveComputeFrom({
+                serverless: connectionManager.serverless,
+                cluster: connectionManager.cluster
+                    ? {id: connectionManager.cluster.id}
+                    : undefined,
+                serverlessVersion: connectionManager.serverlessVersion,
+            });
+            if (resolution.status !== "ok") {
+                return undefined;
+            }
+            try {
+                const result = await pythonSetupClient.run(
+                    {
+                        mode: "default",
+                        dryRun: true,
+                        compute: resolution.compute,
+                    },
+                    {cwd: root, token}
+                );
+                return result.compute?.envKey;
+            } catch {
+                return undefined;
+            }
+        },
+        recordDrift: (report) => telemetry.recordPythonSetupDrift(report),
+    });
+    context.subscriptions.push(
+        pythonSetupDrift,
+        // Compute target changed (cluster attach/detach/switch).
+        connectionManager.onDidChangeCluster(() =>
+            pythonSetupDrift.check("computeChange")
+        ),
+        // Serverless enable/disable and connection churn flow through state.
+        connectionManager.onDidChangeState(() =>
+            pythonSetupDrift.check("computeChange")
+        ),
+        // Re-picking the serverless version while serverless is already selected
+        // fires neither onDidChangeCluster nor onDidChangeState -- it only writes
+        // the `serverlessVersion` config key -- so watch that key directly, or a
+        // v4 -> v2 switch would silently miss drift.
+        configModel.onDidChangeKey("serverlessVersion")(async () =>
+            pythonSetupDrift.check("computeChange")
+        ),
+        // A completed setup updates the persisted state; re-evaluate so a
+        // successful re-run clears the badge promptly.
+        pythonSetupEnvironment.onDidChangeState(() =>
+            pythonSetupDrift.check("setupCompleted")
+        )
+    );
+    // The "workspace open" trigger: evaluate once now that everything is wired.
+    pythonSetupDrift.check("workspaceOpen");
+
+    // The config-view entry combines the setup controller's readiness with the
+    // drift manager's `drifted` signal.
+    const pythonSetupEntry = composePythonSetupEntry(
+        pythonSetupEnvironment,
+        pythonSetupDrift
+    );
+    context.subscriptions.push(pythonSetupEntry);
 
     const environmentCommands = new EnvironmentCommands(
         featureManager,
@@ -1143,7 +1245,7 @@ export async function activate(
         featureManager,
         workspaceFolderManager,
         aiToolsManager,
-        pythonSetupEnvironment
+        pythonSetupEntry
     );
     const configurationView = window.createTreeView("configurationView", {
         treeDataProvider: configurationDataProvider,

packages/databricks-vscode/src/python-setup/controllers/PythonSetupDriftManager.test.ts

@@ -0,0 +1,184 @@
+import {expect} from "chai";
+import {CancellationLike} from "../gateways/PythonSetupCliClient";
+import {
+    PythonSetupDriftDeps,
+    PythonSetupDriftManager,
+} from "./PythonSetupDriftManager";
+
+function makeDeps(over: Partial<PythonSetupDriftDeps> = {}): {
+    deps: PythonSetupDriftDeps;
+    recorded: unknown[];
+    calls: {resolve: number};
+} {
+    const recorded: unknown[] = [];
+    const calls = {resolve: 0};
+    const deps: PythonSetupDriftDeps = {
+        isVisible: async () => true,
+        getPersistedEnvKey: () => "serverless/serverless-v4",
+        getComputeDescriptor: () => "cluster:c1",
+        // eslint-disable-next-line @typescript-eslint/no-unused-vars
+        resolveCurrentEnvKey: async (_token: CancellationLike) => {
+            calls.resolve++;
+            return "dbr/15.4.x-scala2.12";
+        },
+        recordDrift: (r) => recorded.push(r),
+        ...over,
+    };
+    return {deps, recorded, calls};
+}
+
+describe("PythonSetupDriftManager", () => {
+    it("flags drift and reports telemetry when keys differ", async () => {
+        const {deps, recorded} = makeDeps();
+        const m = new PythonSetupDriftManager(deps);
+        let fired = 0;
+        m.onDidChangeState(() => fired++);
+
+        await m.evaluate("computeChange");
+
+        expect(m.drifted).to.be.true;
+        expect(fired).to.equal(1);
+        expect(recorded).to.deep.equal([
+            {
+                trigger: "computeChange",
+                fromEnvKey: "serverless/serverless-v4",
+                toEnvKey: "dbr/15.4.x-scala2.12",
+            },
+        ]);
+        m.dispose();
+    });
+
+    it("does not flag drift when the keys match", async () => {
+        const {deps, recorded} = makeDeps({
+            resolveCurrentEnvKey: async () => "serverless/serverless-v4",
+        });
+        const m = new PythonSetupDriftManager(deps);
+        await m.evaluate("workspaceOpen");
+        expect(m.drifted).to.be.false;
+        expect(recorded).to.have.length(0);
+        m.dispose();
+    });
+
+    it("is a no-op when not visible", async () => {
+        const {deps, recorded} = makeDeps({isVisible: async () => false});
+        const m = new PythonSetupDriftManager(deps);
+        await m.evaluate("workspaceOpen");
+        expect(m.drifted).to.be.false;
+        expect(recorded).to.have.length(0);
+        m.dispose();
+    });
+
+    it("does not flag drift when there is no persisted state", async () => {
+        const {deps} = makeDeps({getPersistedEnvKey: () => undefined});
+        const m = new PythonSetupDriftManager(deps);
+        await m.evaluate("workspaceOpen");
+        expect(m.drifted).to.be.false;
+        m.dispose();
+    });
+
+    it("leaves the flag unchanged when the current key is unknown", async () => {
+        // Start drifted, then a later check can't resolve the key: stay drifted.
+        const {deps} = makeDeps();
+        const m = new PythonSetupDriftManager(deps);
+        await m.evaluate("computeChange");
+        expect(m.drifted).to.be.true;
+
+        (deps as {resolveCurrentEnvKey: unknown}).resolveCurrentEnvKey =
+            async () => undefined;
+        await m.evaluate("workspaceOpen");
+        expect(m.drifted).to.be.true;
+        m.dispose();
+    });
+
+    it("clears drift once the keys match again", async () => {
+        const {deps} = makeDeps();
+        const m = new PythonSetupDriftManager(deps);
+        await m.evaluate("computeChange");
+        expect(m.drifted).to.be.true;
+
+        (deps as {resolveCurrentEnvKey: unknown}).resolveCurrentEnvKey =
+            async () => "serverless/serverless-v4";
+        await m.evaluate("setupCompleted");
+        expect(m.drifted).to.be.false;
+        m.dispose();
+    });
+
+    it("reports the same mismatch only once until it clears", async () => {
+        const {deps, recorded} = makeDeps();
+        const m = new PythonSetupDriftManager(deps);
+        await m.evaluate("computeChange");
+        await m.evaluate("workspaceOpen"); // same mismatch, no new telemetry
+        expect(recorded).to.have.length(1);
+
+        // Clears, then a mismatch recurs on a DIFFERENT compute -> reported
+        // again (a different descriptor is required, since a compute-change with
+        // the same identity is skipped as a no-op runtime-state transition).
+        (deps as {resolveCurrentEnvKey: unknown}).resolveCurrentEnvKey =
+            async () => "serverless/serverless-v4";
+        await m.evaluate("setupCompleted");
+        (deps as {getComputeDescriptor: unknown}).getComputeDescriptor = () =>
+            "cluster:c2";
+        (deps as {resolveCurrentEnvKey: unknown}).resolveCurrentEnvKey =
+            async () => "dbr/15.4.x-scala2.12";
+        await m.evaluate("computeChange");
+        expect(recorded).to.have.length(2);
+        m.dispose();
+    });
+
+    it("skips the dry-run when a compute-change leaves the identity unchanged", async () => {
+        // Same descriptor across two compute-change triggers models a cluster
+        // runtime-state transition (RUNNING -> TERMINATED): the env key cannot
+        // have changed, so the CLI dry-run must not run a second time.
+        const {deps, calls} = makeDeps();
+        const m = new PythonSetupDriftManager(deps);
+        await m.evaluate("computeChange");
+        expect(calls.resolve).to.equal(1);
+        await m.evaluate("computeChange");
+        expect(calls.resolve).to.equal(1);
+        m.dispose();
+    });
+
+    it("clears drift when no comparable compute is attached", async () => {
+        // Start drifted, then compute is detached: drift is meaningless, so the
+        // stale flag must clear rather than linger.
+        const {deps} = makeDeps();
+        const m = new PythonSetupDriftManager(deps);
+        await m.evaluate("computeChange");
+        expect(m.drifted).to.be.true;
+
+        (deps as {getComputeDescriptor: unknown}).getComputeDescriptor = () =>
+            undefined;
+        await m.evaluate("computeChange");
+        expect(m.drifted).to.be.false;
+        m.dispose();
+    });
+
+    it("stays silent and leaves the flag unchanged when a dep rejects", async () => {
+        // A rejecting dep must resolve quietly to "unknown" -- no throw, no
+        // unhandled rejection, and the drift flag is left as-is.
+        const {deps, recorded} = makeDeps({
+            isVisible: async () => {
+                throw new Error("network down");
+            },
+        });
+        const m = new PythonSetupDriftManager(deps);
+
+        // Starts not drifted: a rejection leaves it false.
+        await m.evaluate("workspaceOpen");
+        expect(m.drifted).to.be.false;
+        expect(recorded).to.have.length(0);
+
+        // Now start drifted, then a rejecting dep must not retract the flag.
+        (deps as {isVisible: unknown}).isVisible = async () => true;
+        await m.evaluate("computeChange");
+        expect(m.drifted).to.be.true;
+
+        (deps as {resolveCurrentEnvKey: unknown}).resolveCurrentEnvKey =
+            async () => {
+                throw new Error("dry-run failed");
+            };
+        await m.evaluate("workspaceOpen");
+        expect(m.drifted).to.be.true;
+        m.dispose();
+    });
+});