feat(python-setup): route the setup command to the uv flow when active

packages/databricks-vscode/src/extension.ts

@@ -17,6 +17,7 @@ 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 {routeEnvironmentSetup} from "./language/pythonSetupRouting";
 import {COPY_COMMAND_IDS} from "./ui/configuration-view/copyActions";
 import {AiToolsManager} from "./aitools/AiToolsManager";
 import {AiToolsCommands} from "./aitools/AiToolsCommands";
@@ -911,7 +912,10 @@ export async function activate(
                 pythonExtensionWrapper,
                 environmentDependenciesInstaller,
                 configureAutocomplete,
-                packageManagerTelemetry
+                packageManagerTelemetry,
+                // Constructed lazily (first isEnabled call is well after
+                // pythonSetupEnvironment is wired), so this reference is safe.
+                () => pythonSetupEnvironment.isVisible()
             )
     );
     // uv-native Python environment setup (python-setup). Constructed always,
@@ -1157,8 +1161,16 @@ export async function activate(
     context.subscriptions.push(
         telemetry.registerCommand(
             "databricks.environment.setup",
-            environmentCommands.setup,
-            environmentCommands
+            // Route to the uv-native flow when it is the active surface for the
+            // project, else the legacy checklist. Every trigger surface (status
+            // bar, config-view rows, palette, the run/debug gate) funnels
+            // through this command, so they all dispatch here.
+            (stepId?: string) =>
+                routeEnvironmentSetup(
+                    pythonSetupEnvironment,
+                    environmentCommands,
+                    stepId
+                )
         ),
         telemetry.registerCommand(
             "databricks.environment.refresh",

packages/databricks-vscode/src/language/EnvironmentDependenciesVerifier.test.ts

@@ -0,0 +1,86 @@
+import * as assert from "assert";
+import type {Disposable} from "vscode";
+import {EnvironmentDependenciesVerifier} from "./EnvironmentDependenciesVerifier";
+import {ConnectionManager} from "../configuration/ConnectionManager";
+import {MsPythonExtensionWrapper} from "./MsPythonExtensionWrapper";
+import {EnvironmentDependenciesInstaller} from "./EnvironmentDependenciesInstaller";
+import {ConfigureAutocomplete} from "./ConfigureAutocomplete";
+import {PackageManagerTelemetry} from "./PackageManagerTelemetry";
+
+// ConnectionManager/MsPythonExtensionWrapper expose their VS Code Events as
+// instance properties (`emitter.event`), which ts-mockito can't stub — calling
+// them throws "not a function" during construction. Hand-rolled stubs give the
+// constructor real no-op event subscriptions and let us drive the one handler
+// under test directly.
+const noopEvent = () => ({dispose() {}}) as Disposable;
+
+describe(__filename, () => {
+    let showCalls: unknown[];
+
+    function makeVerifier(isUvActive: () => Promise<boolean>) {
+        const connectionManager = {
+            serverless: false,
+            cluster: undefined,
+            onDidChangeCluster: noopEvent,
+            onDidChangeState: noopEvent,
+        } as unknown as ConnectionManager;
+
+        const pythonExtension = {
+            pythonEnvironment: Promise.resolve({
+                version: {major: 3, minor: 10, micro: 0},
+                environment: {name: ".venv"},
+                executable: {uri: {fsPath: "/project/.venv/bin/python"}},
+            }),
+            getPythonExecutable: async () => "/project/.venv/bin/python",
+            // databricks-connect missing: the legacy path would offer to install
+            // it on an interpreter change.
+            getPackageDetailsFromEnvironment: async () => undefined,
+            onDidChangePythonExecutable: noopEvent,
+        } as unknown as MsPythonExtensionWrapper;
+
+        const installer = {
+            show: (advertisement?: boolean) => {
+                showCalls.push(advertisement);
+                return Promise.resolve();
+            },
+            onDidTryInstallation: noopEvent,
+        } as unknown as EnvironmentDependenciesInstaller;
+
+        const configureAutocomplete = {
+            shouldSetupBuiltins: async () => false,
+            onDidUpdate: noopEvent,
+        } as unknown as ConfigureAutocomplete;
+
+        const packageManagerTelemetry =
+            {} as unknown as PackageManagerTelemetry;
+
+        return new EnvironmentDependenciesVerifier(
+            connectionManager,
+            pythonExtension,
+            installer,
+            configureAutocomplete,
+            packageManagerTelemetry,
+            isUvActive
+        );
+    }
+
+    beforeEach(() => {
+        showCalls = [];
+    });
+
+    it("offers the legacy install prompt on interpreter change when uv is not active", async () => {
+        const verifier = makeVerifier(async () => false);
+
+        await verifier["onInterpreterChanged"]();
+
+        assert.deepStrictEqual(showCalls, [true]);
+    });
+
+    it("suppresses the legacy install prompt on interpreter change when uv is active", async () => {
+        const verifier = makeVerifier(async () => true);
+
+        await verifier["onInterpreterChanged"]();
+
+        assert.deepStrictEqual(showCalls, []);
+    });
+});

packages/databricks-vscode/src/language/EnvironmentDependenciesVerifier.ts

@@ -20,7 +20,12 @@ export class EnvironmentDependenciesVerifier extends MultiStepAccessVerifier {
         private readonly pythonExtension: MsPythonExtensionWrapper,
         private readonly installer: EnvironmentDependenciesInstaller,
         private readonly configureAutocomplete: ConfigureAutocomplete,
-        private readonly packageManagerTelemetry: PackageManagerTelemetry
+        private readonly packageManagerTelemetry: PackageManagerTelemetry,
+        // Whether the uv-native flow is the active surface for this project. When
+        // it is, this legacy checklist stays evaluated (other consumers read its
+        // state) but must not auto-prompt, or the user gets both flows' prompts.
+        // Defaults to "never active" so the legacy construction is unchanged.
+        private readonly isUvActive: () => Promise<boolean> = async () => false
     ) {
         super([
             "checkCluster",
@@ -43,14 +48,10 @@ export class EnvironmentDependenciesVerifier extends MultiStepAccessVerifier {
                     await this.checkWorkspaceHasUc();
                 }
             }, this),
-            this.pythonExtension.onDidChangePythonExecutable(async () => {
-                await this.checkPythonEnvironment();
-                const depsCheck = await this.checkEnvironmentDependencies();
-                if (!depsCheck.available && depsCheck.action) {
-                    await depsCheck.action(true);
-                }
-                await this.checkBuiltins();
-            }, this),
+            this.pythonExtension.onDidChangePythonExecutable(
+                () => this.onInterpreterChanged(),
+                this
+            ),
             this.installer.onDidTryInstallation(async () => {
                 await this.checkEnvironmentDependencies();
                 await this.checkBuiltins();
@@ -61,6 +62,22 @@ export class EnvironmentDependenciesVerifier extends MultiStepAccessVerifier {
         );
     }
 
+    private async onInterpreterChanged() {
+        await this.checkPythonEnvironment();
+        const depsCheck = await this.checkEnvironmentDependencies();
+        // Suppress the legacy auto-install prompt when the uv-native flow owns
+        // this project — it drives its own setup, so a second prompt here is a
+        // duplicate.
+        if (
+            !depsCheck.available &&
+            depsCheck.action &&
+            !(await this.isUvActive())
+        ) {
+            await depsCheck.action(true);
+        }
+        await this.checkBuiltins();
+    }
+
     promptForAttachingCluster(msg: string) {
         return async () => {
             await commands.executeCommand(

packages/databricks-vscode/src/language/pythonSetupRouting.test.ts

@@ -0,0 +1,47 @@
+import * as assert from "assert";
+import {routeEnvironmentSetup} from "./pythonSetupRouting";
+
+describe(__filename, () => {
+    function makeUv(visible: boolean) {
+        const calls = {setup: 0};
+        return {
+            uv: {
+                isVisible: () => Promise.resolve(visible),
+                setup: async () => {
+                    calls.setup++;
+                },
+            },
+            calls,
+        };
+    }
+
+    it("routes to the uv flow when it is the active surface", async () => {
+        const {uv, calls} = makeUv(true);
+        const legacyCalls: (string | undefined)[] = [];
+        const legacy = {
+            setup: async (stepId?: string) => {
+                legacyCalls.push(stepId);
+            },
+        };
+
+        await routeEnvironmentSetup(uv, legacy, "checkPythonEnvironment");
+
+        assert.strictEqual(calls.setup, 1);
+        assert.deepStrictEqual(legacyCalls, []);
+    });
+
+    it("routes to the legacy checklist (with the step id) when uv is not active", async () => {
+        const {uv, calls} = makeUv(false);
+        const legacyCalls: (string | undefined)[] = [];
+        const legacy = {
+            setup: async (stepId?: string) => {
+                legacyCalls.push(stepId);
+            },
+        };
+
+        await routeEnvironmentSetup(uv, legacy, "checkPythonEnvironment");
+
+        assert.strictEqual(calls.setup, 0);
+        assert.deepStrictEqual(legacyCalls, ["checkPythonEnvironment"]);
+    });
+});

packages/databricks-vscode/src/language/pythonSetupRouting.ts

@@ -0,0 +1,38 @@
+/**
+ * The slice of the uv-native setup flow ({@link
+ * ../python-setup/controllers/PythonSetupEnvironmentSetup}) that the setup
+ * command router needs. Kept as a narrow structural interface so the router
+ * carries no dependency on the controller/gateway layers and stays
+ * unit-testable.
+ */
+export interface UvPythonSetup {
+    /** Whether the uv-native flow is the active surface for the current project. */
+    isVisible(): Promise<boolean>;
+    /** Run the uv-native setup (re-entrancy-guarded); resolves when it settles. */
+    setup(): Promise<void>;
+}
+
+/** The legacy checklist setup, as the router invokes it. */
+export interface LegacyEnvironmentSetup {
+    setup(stepId?: string): Promise<void>;
+}
+
+/**
+ * Back the `databricks.environment.setup` command: run the uv-native flow when
+ * it is the active surface for the current project, otherwise the legacy
+ * checklist. Routing here means every surface that funnels through that command
+ * (status bar, config-view rows, palette, the run/debug gate) reaches the right
+ * flow without per-surface branching. `stepId` is a legacy-only affordance and
+ * is ignored by the uv flow, which has no per-step entry points.
+ */
+export async function routeEnvironmentSetup(
+    uv: UvPythonSetup,
+    legacy: LegacyEnvironmentSetup,
+    stepId?: string
+): Promise<void> {
+    if (await uv.isVisible()) {
+        await uv.setup();
+        return;
+    }
+    await legacy.setup(stepId);
+}

packages/databricks-vscode/src/run/RunCommands.test.ts

@@ -1,6 +1,6 @@
 import * as assert from "assert";
 import {commands, ExtensionContext, Uri} from "vscode";
-import {anything, instance, mock, verify, when} from "ts-mockito";
+import {instance, mock, verify, when} from "ts-mockito";
 import {RunCommands} from "./RunCommands";
 import {ConnectionManager} from "../configuration/ConnectionManager";
 import {MsPythonExtensionWrapper} from "../language/MsPythonExtensionWrapper";
@@ -119,10 +119,16 @@ describe(__filename, () => {
             ).never();
         });
 
-        it("should not re-verify when the feature is unavailable", async () => {
-            when(featureManagerMock.isEnabled(anything())).thenResolve(
-                featureState(false, undefined)
-            );
+        it("should not run the staleness re-check when the feature is unavailable", async () => {
+            // Unavailable on the first (cached) check, so the staleness re-check
+            // is skipped; setup runs and the forced post-setup check still
+            // reports unavailable, so the launch aborts.
+            when(
+                featureManagerMock.isEnabled("environment.dependencies")
+            ).thenResolve(featureState(false, undefined));
+            when(
+                featureManagerMock.isEnabled("environment.dependencies", true)
+            ).thenResolve(featureState(false, undefined));
             when(pythonExtensionMock.getPythonExecutable()).thenResolve(
                 undefined
             );
@@ -134,26 +140,21 @@ describe(__filename, () => {
             try {
                 const result = await runCommands["checkDbconnectEnabled"]();
                 assert.strictEqual(result, false);
-                verify(
-                    featureManagerMock.isEnabled(
-                        "environment.dependencies",
-                        true
-                    )
-                ).never();
             } finally {
                 registration.dispose();
             }
         });
 
         it("should proceed after a successful in-flow setup", async () => {
             // Unavailable on the first check; the setup flow fixes it, so the
-            // re-check after setup reports available and the launch proceeds.
+            // forced re-check after setup reports available and the launch
+            // proceeds.
             when(
                 featureManagerMock.isEnabled("environment.dependencies")
-            ).thenResolve(
-                featureState(false, undefined),
-                featureState(true, "/project/.venv/bin/python")
-            );
+            ).thenResolve(featureState(false, undefined));
+            when(
+                featureManagerMock.isEnabled("environment.dependencies", true)
+            ).thenResolve(featureState(true, "/project/.venv/bin/python"));
             const registration = commands.registerCommand(
                 "databricks.environment.setup",
                 () => undefined

packages/databricks-vscode/src/run/RunCommands.ts

@@ -138,9 +138,17 @@ export class RunCommands {
         }
         // Run the setup flow, then re-check: a successful setup should let the
         // launch proceed instead of aborting and making the user re-trigger.
+        // Force a fresh check rather than reading the cache: when the setup
+        // command routes to the uv flow, it adopts the interpreter and the
+        // legacy state refreshes only on the (async) interpreter-change event,
+        // which may not have landed yet.
         await commands.executeCommand("databricks.environment.setup");
-        return (await this.featureManager.isEnabled("environment.dependencies"))
-            .available;
+        return (
+            await this.featureManager.isEnabled(
+                "environment.dependencies",
+                true
+            )
+        ).available;
     }
 
     private async isPythonEnvironmentStale(featureState: FeatureState) {