Add AI tools integration

packages/databricks-vscode/package.json

@@ -91,6 +91,48 @@
                 "enablement": "databricks.context.activated && databricks.context.loggedIn && !databricks.context.remoteMode",
                 "icon": "$(debug-disconnect)"
             },
+            {
+                "command": "databricks.aitools.install",
+                "title": "Install AI tools",
+                "category": "Databricks",
+                "icon": "$(cloud-download)"
+            },
+            {
+                "command": "databricks.aitools.checkForUpdates",
+                "title": "Check for AI tools updates",
+                "category": "Databricks",
+                "icon": "$(refresh)"
+            },
+            {
+                "command": "databricks.aitools.reload",
+                "title": "Reload AI tools",
+                "category": "Databricks",
+                "icon": "$(refresh)"
+            },
+            {
+                "command": "databricks.aitools.update",
+                "title": "Update AI tools",
+                "category": "Databricks",
+                "icon": "$(arrow-circle-up)"
+            },
+            {
+                "command": "databricks.aitools.uninstall",
+                "title": "Uninstall AI tools",
+                "category": "Databricks",
+                "icon": "$(trash)"
+            },
+            {
+                "command": "databricks.aitools.addCursorPlugin",
+                "title": "Add Databricks plugin to Cursor",
+                "category": "Databricks",
+                "icon": "$(plug)"
+            },
+            {
+                "command": "databricks.developer.resetState",
+                "title": "Reset State (Developer)",
+                "category": "Databricks",
+                "icon": "$(debug-restart)"
+            },
             {
                 "command": "databricks.cluster.filterByAll",
                 "title": "All",
@@ -986,6 +1028,16 @@
                     "when": "view == configurationView && viewItem =~ /databricks.configuration.cluster.*\\.terminated.*/ && databricks.context.bundle.deploymentState == idle",
                     "group": "navigation@0"
                 },
+                {
+                    "command": "databricks.aitools.addCursorPlugin",
+                    "when": "view == configurationView && viewItem =~ /^databricks.configuration.aitools.(installed|upToDate|updateAvailable|checking|updating|error)$/ && databricks.context.aitools.showCursorPlugin",
+                    "group": "inline@0"
+                },
+                {
+                    "command": "databricks.aitools.uninstall",
+                    "when": "view == configurationView && viewItem =~ /^databricks.configuration.aitools.(installed|upToDate|updateAvailable|checking|updating|error)$/",
+                    "group": "navigation@9"
+                },
                 {
                     "command": "databricks.bundle.deployAndRunJob",
                     "when": "view == dabsResourceExplorerView && viewItem =~ /^databricks.bundle.resource=jobs.runnable.*$/ && databricks.context.bundle.deploymentState == idle",
@@ -1138,6 +1190,34 @@
                 }
             ],
             "commandPalette": [
+                {
+                    "command": "databricks.developer.resetState",
+                    "when": "databricks.context.development"
+                },
+                {
+                    "command": "databricks.aitools.addCursorPlugin",
+                    "when": "databricks.context.aitools.showCursorPlugin && !databricks.context.remoteMode"
+                },
+                {
+                    "command": "databricks.aitools.install",
+                    "when": "!databricks.context.aitools.installed && !databricks.context.remoteMode"
+                },
+                {
+                    "command": "databricks.aitools.uninstall",
+                    "when": "databricks.context.aitools.installed && !databricks.context.remoteMode"
+                },
+                {
+                    "command": "databricks.aitools.update",
+                    "when": "databricks.context.aitools.installed && !databricks.context.remoteMode"
+                },
+                {
+                    "command": "databricks.aitools.checkForUpdates",
+                    "when": "databricks.context.aitools.installed && !databricks.context.remoteMode"
+                },
+                {
+                    "command": "databricks.aitools.reload",
+                    "when": "false"
+                },
                 {
                     "command": "databricks.run.runEditorContents",
                     "when": "resourceLangId == python"

packages/databricks-vscode/src/aitools/AiToolsCommands.ts

@@ -0,0 +1,129 @@
+import {Disposable, QuickPickItem, window} from "vscode";
+import {AiToolsManager} from "./AiToolsManager";
+import {AiToolsScope} from "../cli/CliWrapper";
+import {AiToolsInstallSource} from "../telemetry/constants";
+
+interface ScopeQuickPickItem extends QuickPickItem {
+    scope: AiToolsScope;
+}
+
+export class AiToolsCommands implements Disposable {
+    private disposables: Disposable[] = [];
+
+    constructor(private readonly aiToolsManager: AiToolsManager) {}
+
+    dispose() {
+        this.disposables.forEach((d) => d.dispose());
+    }
+
+    installCommand() {
+        // The command may be invoked with a source argument (e.g. the first-load
+        // modal passes "modal"); default to "sidePane" for the manual affordance.
+        return async (source: AiToolsInstallSource = "sidePane") => {
+            const scope = await this.pickScope();
+            if (scope === undefined) {
+                return;
+            }
+            await this.aiToolsManager.install(scope, source);
+        };
+    }
+
+    /**
+     * Show the scope picker. Project scope is always listed, but is shown
+     * disabled (greyed, with a hint) and cannot be selected when no workspace
+     * folder is open, since it installs into the open folder. Resolves to the
+     * chosen scope, or undefined if the picker was dismissed.
+     */
+    private pickScope(): Promise<AiToolsScope | undefined> {
+        const hasFolder = this.aiToolsManager.hasProjectFolder;
+        const quickPick = window.createQuickPick<ScopeQuickPickItem>();
+        quickPick.title = "Install Databricks AI tools";
+        quickPick.placeholder = "Choose where to install the AI tools";
+        quickPick.items = [
+            {
+                label: "$(globe) Global",
+                detail: "Install AI tools for all projects on this machine",
+                scope: "global",
+            },
+            {
+                label: "$(folder) Project",
+                detail: hasFolder
+                    ? "Install AI tools into this project"
+                    : "Open a folder to install AI tools into a project",
+                // Render as disabled when there's no folder to install into.
+                description: hasFolder ? undefined : "Requires an open folder",
+                scope: "project",
+            },
+        ];
+
+        return new Promise<AiToolsScope | undefined>((resolve) => {
+            let resolved: AiToolsScope | undefined;
+            this.disposables.push(
+                quickPick.onDidAccept(() => {
+                    const picked = quickPick.selectedItems[0];
+                    // Ignore selection of the disabled project item; keep the
+                    // picker open so the choice reads as non-actionable.
+                    if (
+                        picked === undefined ||
+                        (picked.scope === "project" && !hasFolder)
+                    ) {
+                        return;
+                    }
+                    resolved = picked.scope;
+                    quickPick.hide();
+                }),
+                quickPick.onDidHide(() => {
+                    resolve(resolved);
+                    quickPick.dispose();
+                })
+            );
+            quickPick.show();
+        });
+    }
+
+    checkForUpdatesCommand() {
+        return async () => {
+            await this.aiToolsManager.checkForUpdates();
+        };
+    }
+
+    /**
+     * Re-run install detection (and update check if installed). Used by the
+     * error row to recover after a transient detection failure.
+     */
+    reloadCommand() {
+        return async () => {
+            await this.aiToolsManager.initialize();
+        };
+    }
+
+    updateCommand() {
+        return async () => {
+            await this.aiToolsManager.update();
+        };
+    }
+
+    uninstallCommand() {
+        return async () => {
+            const location = this.aiToolsManager.state.installLocation;
+            if (location === undefined) {
+                return;
+            }
+            const confirm = await window.showWarningMessage(
+                `Uninstall Databricks AI tools (${location})?`,
+                {modal: true},
+                "Uninstall"
+            );
+            if (confirm !== "Uninstall") {
+                return;
+            }
+            await this.aiToolsManager.uninstall();
+        };
+    }
+
+    addCursorPluginCommand() {
+        return async () => {
+            await this.aiToolsManager.addCursorPlugin();
+        };
+    }
+}