diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index 682c9a29c..910a8a7b6 100644 --- a/packages/databricks-vscode/package.json +++ b/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" diff --git a/packages/databricks-vscode/src/aitools/AiToolsCommands.ts b/packages/databricks-vscode/src/aitools/AiToolsCommands.ts new file mode 100644 index 000000000..51f50adc4 --- /dev/null +++ b/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 { + const hasFolder = this.aiToolsManager.hasProjectFolder; + const quickPick = window.createQuickPick(); + 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((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(); + }; + } +} diff --git a/packages/databricks-vscode/src/aitools/AiToolsManager.test.ts b/packages/databricks-vscode/src/aitools/AiToolsManager.test.ts new file mode 100644 index 000000000..fc91c04e0 --- /dev/null +++ b/packages/databricks-vscode/src/aitools/AiToolsManager.test.ts @@ -0,0 +1,658 @@ +/* eslint-disable @typescript-eslint/naming-convention */ + +import assert from "assert"; +import {anything, instance, mock, reset, verify, when} from "ts-mockito"; +import {commands, Uri, window} from "vscode"; +import {mkdtemp, mkdir, writeFile, rm} from "fs/promises"; +import path from "path"; +import os from "os"; +import {AiToolsListResult, CliWrapper, ProcessError} from "../cli/CliWrapper"; +import {StateStorage} from "../vscode-objs/StateStorage"; +import {WorkspaceFolderManager} from "../vscode-objs/WorkspaceFolderManager"; +import {Telemetry} from "../telemetry"; +import {AiToolsManager} from "./AiToolsManager"; + +const STATE_FILE_RELATIVE_PATH = path.join( + ".databricks", + "aitools", + "skills", + ".state.json" +); + +function listResult( + skills: Array<{ + name: string; + latest_version: string; + installed: Record; + }> +): AiToolsListResult { + return { + release: "0.2.9", + skills: skills.map((s) => ({ + experimental: false, + ...s, + })), + }; +} + +describe(__filename, () => { + let mockCli: CliWrapper; + let mockWorkspaceFolderManager: WorkspaceFolderManager; + let telemetry: Telemetry; + let storedState: Record; + let stubStateStorage: StateStorage; + + let projectDir: string; + let homeDir: string; + let originalHome: string | undefined; + + async function writeStateFile(root: string) { + const dir = path.join(root, path.dirname(STATE_FILE_RELATIVE_PATH)); + await mkdir(dir, {recursive: true}); + await writeFile( + path.join(root, STATE_FILE_RELATIVE_PATH), + JSON.stringify({schema_version: 1, release: "v0.2.9", skills: {}}) + ); + } + + beforeEach(async () => { + projectDir = await mkdtemp(path.join(os.tmpdir(), "aitools-proj-")); + homeDir = await mkdtemp(path.join(os.tmpdir(), "aitools-home-")); + originalHome = process.env.HOME; + process.env.HOME = homeDir; + + storedState = {}; + stubStateStorage = { + get: (key: string) => storedState[key], + set: async (key: string, value: any) => { + storedState[key] = value; + }, + onDidChange: () => ({dispose() {}}), + } as unknown as StateStorage; + + mockCli = mock(CliWrapper); + mockWorkspaceFolderManager = mock(WorkspaceFolderManager); + when(mockWorkspaceFolderManager.activeProjectUri).thenReturn( + Uri.file(projectDir) + ); + // start() returns a recorder callback; stub it to a no-op so the + // manager's install/update/uninstall telemetry calls work. + telemetry = { + start: () => () => {}, + } as unknown as Telemetry; + }); + + afterEach(async () => { + process.env.HOME = originalHome; + reset(mockCli); + reset(mockWorkspaceFolderManager); + await rm(projectDir, {recursive: true, force: true}); + await rm(homeDir, {recursive: true, force: true}); + }); + + function createManager() { + return new AiToolsManager( + instance(mockCli), + stubStateStorage, + instance(mockWorkspaceFolderManager), + telemetry + ); + } + + it("detects no install when no state file exists", async () => { + const manager = createManager(); + const location = await manager.detectInstall(); + assert.strictEqual(location, undefined); + assert.strictEqual(manager.isInstalled, false); + assert.strictEqual( + storedState["databricks.aitools.installLocation"], + undefined + ); + }); + + it("detects a project install", async () => { + await writeStateFile(projectDir); + const manager = createManager(); + const location = await manager.detectInstall(); + assert.strictEqual(location, "project"); + assert.strictEqual(manager.isInstalled, true); + assert.strictEqual( + storedState["databricks.aitools.installLocation"], + "project" + ); + }); + + it("detects a global install when only the home state file exists", async () => { + await writeStateFile(homeDir); + const manager = createManager(); + const location = await manager.detectInstall(); + assert.strictEqual(location, "global"); + assert.strictEqual( + storedState["databricks.aitools.installLocation"], + "global" + ); + }); + + it("prefers project over global when both exist", async () => { + await writeStateFile(projectDir); + await writeStateFile(homeDir); + const manager = createManager(); + assert.strictEqual(await manager.detectInstall(), "project"); + }); + + it("preserves the cached location on an unexpected detection error", async () => { + // First, a clean detect that finds a project install. + await writeStateFile(projectDir); + const manager = createManager(); + assert.strictEqual(await manager.detectInstall(), "project"); + assert.strictEqual(manager.state.detectError ?? false, false); + + // Now make the state file unreadable as a file: replace it with a + // directory so readFile throws EISDIR (a non-ENOENT error). + await rm(path.join(projectDir, STATE_FILE_RELATIVE_PATH), { + force: true, + }); + await mkdir(path.join(projectDir, STATE_FILE_RELATIVE_PATH)); + + const location = await manager.detectInstall(); + + // Location is preserved (not flipped to undefined) and the error flag is set. + assert.strictEqual(location, "project"); + assert.strictEqual(manager.state.installLocation, "project"); + assert.strictEqual(manager.state.detectError, true); + assert.strictEqual( + storedState["databricks.aitools.installLocation"], + "project" + ); + }); + + it("clears the detect error flag on a subsequent successful detect", async () => { + await writeStateFile(projectDir); + const manager = createManager(); + await manager.detectInstall(); + + // Trigger an error (state file is a directory), then recover. + await rm(path.join(projectDir, STATE_FILE_RELATIVE_PATH), { + force: true, + }); + await mkdir(path.join(projectDir, STATE_FILE_RELATIVE_PATH)); + await manager.detectInstall(); + assert.strictEqual(manager.state.detectError, true); + + // Restore a real state file; detection should succeed and clear the flag. + await rm(path.join(projectDir, STATE_FILE_RELATIVE_PATH), { + recursive: true, + force: true, + }); + await writeStateFile(projectDir); + await manager.detectInstall(); + assert.strictEqual(manager.state.detectError, false); + assert.strictEqual(manager.state.installLocation, "project"); + }); + + it("reports upToDate when all installed skills match latest", async () => { + await writeStateFile(projectDir); + when(mockCli.aitoolsList(anything())).thenResolve( + listResult([ + { + name: "databricks-core", + latest_version: "0.1.0", + installed: {project: "0.1.0"}, + }, + ]) + ); + const manager = createManager(); + await manager.detectInstall(); + const status = await manager.checkForUpdates(); + assert.strictEqual(status, "upToDate"); + assert.strictEqual(manager.state.updateStatus, "upToDate"); + }); + + it("reports updateAvailable when an installed skill is behind latest", async () => { + await writeStateFile(projectDir); + when(mockCli.aitoolsList(anything())).thenResolve( + listResult([ + { + name: "databricks-core", + latest_version: "0.1.0", + installed: {project: "0.0.1"}, + }, + { + name: "databricks-jobs", + latest_version: "0.2.0", + installed: {project: "0.2.0"}, + }, + ]) + ); + const manager = createManager(); + await manager.detectInstall(); + assert.strictEqual(await manager.checkForUpdates(), "updateAvailable"); + }); + + it("ignores non-installed skills when computing update status", async () => { + await writeStateFile(projectDir); + when(mockCli.aitoolsList(anything())).thenResolve( + listResult([ + { + name: "databricks-core", + latest_version: "0.1.0", + installed: {project: "0.1.0"}, + }, + { + // Not installed (empty installed map) -> must not count as + // an available update even though latest > "". + name: "databricks-uninstalled", + latest_version: "9.9.9", + installed: {}, + }, + ]) + ); + const manager = createManager(); + await manager.detectInstall(); + assert.strictEqual(await manager.checkForUpdates(), "upToDate"); + }); + + it("reports error when the list command fails", async () => { + await writeStateFile(projectDir); + when(mockCli.aitoolsList(anything())).thenReject(new Error("boom")); + const manager = createManager(); + await manager.detectInstall(); + assert.strictEqual(await manager.checkForUpdates(), "error"); + }); + + it("returns unknown update status when not installed", async () => { + const manager = createManager(); + await manager.detectInstall(); + const status = await manager.checkForUpdates(); + assert.strictEqual(status, "unknown"); + verify(mockCli.aitoolsList(anything())).never(); + }); + + it("uninstalls for the detected scope and re-detects", async () => { + await writeStateFile(projectDir); + when( + mockCli.aitoolsUninstall("project", anything(), anything()) + ).thenCall(async () => { + await rm(path.join(projectDir, STATE_FILE_RELATIVE_PATH), { + force: true, + }); + }); + const manager = createManager(); + await manager.detectInstall(); + assert.strictEqual(manager.isInstalled, true); + + await manager.uninstall(); + + verify( + mockCli.aitoolsUninstall("project", anything(), anything()) + ).once(); + assert.strictEqual(manager.isInstalled, false); + assert.strictEqual( + storedState["databricks.aitools.installLocation"], + undefined + ); + }); + + it("toggles the installed when-context on detect and uninstall", async () => { + const contextValues: Array = []; + const original = commands.executeCommand; + (commands as any).executeCommand = ( + command: string, + ...args: any[] + ) => { + if ( + command === "setContext" && + args[0] === "databricks.context.aitools.installed" + ) { + contextValues.push(args[1]); + } + }; + try { + await writeStateFile(projectDir); + when( + mockCli.aitoolsUninstall("project", anything(), anything()) + ).thenCall(async () => { + await rm(path.join(projectDir, STATE_FILE_RELATIVE_PATH), { + force: true, + }); + }); + const manager = createManager(); + + await manager.detectInstall(); + await manager.uninstall(); + + // Last value must reflect "not installed" after uninstall. + assert.strictEqual(contextValues.at(-1), false); + // And it was true at some point (after detecting the install). + assert.ok(contextValues.includes(true)); + } finally { + (commands as any).executeCommand = original; + } + }); + + it("clears the Cursor plugin flag on uninstall so it is re-offered", async () => { + await writeStateFile(projectDir); + storedState["databricks.aitools.cursorPluginPrompted"] = true; + when( + mockCli.aitoolsUninstall("project", anything(), anything()) + ).thenCall(async () => { + await rm(path.join(projectDir, STATE_FILE_RELATIVE_PATH), { + force: true, + }); + }); + const manager = createManager(); + await manager.detectInstall(); + + await manager.uninstall(); + + assert.strictEqual( + storedState["databricks.aitools.cursorPluginPrompted"], + false + ); + }); + + it("does not clear the Cursor plugin flag when uninstall fails", async () => { + await writeStateFile(projectDir); + storedState["databricks.aitools.cursorPluginPrompted"] = true; + when( + mockCli.aitoolsUninstall("project", anything(), anything()) + ).thenReject(new ProcessError("boom", 1)); + const manager = createManager(); + await manager.detectInstall(); + + await manager.uninstall(); + + // The state file still exists (uninstall failed), so the flag must be + // left untouched. + assert.strictEqual( + storedState["databricks.aitools.cursorPluginPrompted"], + true + ); + }); + + it("does not call the CLI when uninstalling with nothing installed", async () => { + const manager = createManager(); + await manager.detectInstall(); + await manager.uninstall(); + verify( + mockCli.aitoolsUninstall(anything(), anything(), anything()) + ).never(); + }); + + it("detects install and refreshes status after a global install", async () => { + when(mockCli.aitoolsInstall("global", anything(), anything())).thenCall( + async () => { + await writeStateFile(homeDir); + } + ); + when(mockCli.aitoolsList(anything())).thenResolve( + listResult([ + { + name: "databricks-core", + latest_version: "0.1.0", + installed: {global: "0.1.0"}, + }, + ]) + ); + const manager = createManager(); + + await manager.install("global"); + + verify(mockCli.aitoolsInstall("global", anything(), anything())).once(); + assert.strictEqual(manager.state.installLocation, "global"); + assert.strictEqual(manager.state.updateStatus, "upToDate"); + }); + + it("refreshes update status to upToDate after a successful update", async () => { + await writeStateFile(projectDir); + when( + mockCli.aitoolsUpdate("project", anything(), anything()) + ).thenResolve(); + // After the update, list reports everything at the latest version. + when(mockCli.aitoolsList(anything())).thenResolve( + listResult([ + { + name: "databricks-core", + latest_version: "0.1.0", + installed: {project: "0.1.0"}, + }, + ]) + ); + const manager = createManager(); + await manager.detectInstall(); + + await manager.update(); + + assert.strictEqual(manager.state.updateStatus, "upToDate"); + verify(mockCli.aitoolsList(anything())).once(); + }); + + it("still refreshes update status when the update command fails", async () => { + await writeStateFile(projectDir); + when( + mockCli.aitoolsUpdate("project", anything(), anything()) + ).thenReject(new ProcessError("boom", 1)); + when(mockCli.aitoolsList(anything())).thenResolve( + listResult([ + { + name: "databricks-core", + latest_version: "0.1.0", + installed: {project: "0.0.1"}, + }, + ]) + ); + const manager = createManager(); + await manager.detectInstall(); + + await manager.update(); + + // The finally block reconciles state even though the update errored. + assert.strictEqual(manager.state.updateStatus, "updateAvailable"); + verify(mockCli.aitoolsList(anything())).once(); + }); + + it("captures the installed release version from list", async () => { + await writeStateFile(projectDir); + when(mockCli.aitoolsList(anything())).thenResolve({ + release: "0.3.1", + skills: [ + { + name: "databricks-core", + latest_version: "0.1.0", + experimental: false, + installed: {project: "0.1.0"}, + }, + ], + }); + const manager = createManager(); + await manager.detectInstall(); + await manager.checkForUpdates(); + assert.strictEqual(manager.state.version, "0.3.1"); + }); + + it("clears the version when nothing is installed", async () => { + await writeStateFile(projectDir); + when(mockCli.aitoolsList(anything())).thenResolve( + listResult([ + { + name: "databricks-core", + latest_version: "0.1.0", + installed: {project: "0.1.0"}, + }, + ]) + ); + const manager = createManager(); + await manager.detectInstall(); + await manager.checkForUpdates(); + assert.strictEqual(manager.state.version, "0.2.9"); + + // Uninstalling / re-detecting with no state file clears the version. + await rm(path.join(projectDir, STATE_FILE_RELATIVE_PATH), { + force: true, + }); + await manager.detectInstall(); + assert.strictEqual(manager.state.version, undefined); + }); + + describe("no open folder", () => { + beforeEach(() => { + // Mirror WorkspaceFolderManager throwing when no folder is active. + when(mockWorkspaceFolderManager.activeProjectUri).thenThrow( + new Error("No active project folder") + ); + }); + + it("reports no project folder", () => { + assert.strictEqual(createManager().hasProjectFolder, false); + }); + + it("installs global against the home dir without touching projectRoot", async () => { + when( + mockCli.aitoolsInstall("global", anything(), anything()) + ).thenCall(async () => { + await writeStateFile(homeDir); + }); + when(mockCli.aitoolsList(anything())).thenResolve( + listResult([ + { + name: "databricks-core", + latest_version: "0.1.0", + installed: {global: "0.1.0"}, + }, + ]) + ); + const manager = createManager(); + + // Must not throw even though no folder is open. + await manager.install("global"); + + verify( + mockCli.aitoolsInstall("global", homeDir, anything()) + ).once(); + assert.strictEqual(manager.state.installLocation, "global"); + }); + + it("detects a global install with no folder open", async () => { + await writeStateFile(homeDir); + const manager = createManager(); + assert.strictEqual(await manager.detectInstall(), "global"); + }); + }); + + describe("initialize", () => { + let originalExecuteCommand: typeof commands.executeCommand; + let originalShowInfo: typeof window.showInformationMessage; + let executed: Array<{command: string; args: any[]}>; + + beforeEach(() => { + originalExecuteCommand = commands.executeCommand; + originalShowInfo = window.showInformationMessage; + executed = []; + (commands as any).executeCommand = async ( + command: string, + ...args: any[] + ) => { + executed.push({command, args}); + }; + }); + + afterEach(() => { + (commands as any).executeCommand = originalExecuteCommand; + (window as any).showInformationMessage = originalShowInfo; + }); + + it("auto-applies an available update when installed", async () => { + await writeStateFile(projectDir); + when( + mockCli.aitoolsUpdate("project", anything(), anything()) + ).thenResolve(); + let call = 0; + when(mockCli.aitoolsList(anything())).thenCall(async () => { + call++; + // First check: behind latest. After update: up to date. + return listResult([ + { + name: "databricks-core", + latest_version: "0.1.0", + installed: {project: call === 1 ? "0.0.1" : "0.1.0"}, + }, + ]); + }); + const manager = createManager(); + + await manager.initialize(); + + verify( + mockCli.aitoolsUpdate("project", anything(), anything()) + ).once(); + assert.strictEqual(manager.state.updateStatus, "upToDate"); + }); + + it("does not update when already up to date", async () => { + await writeStateFile(projectDir); + when(mockCli.aitoolsList(anything())).thenResolve( + listResult([ + { + name: "databricks-core", + latest_version: "0.1.0", + installed: {project: "0.1.0"}, + }, + ]) + ); + const manager = createManager(); + + await manager.initialize(); + + verify( + mockCli.aitoolsUpdate(anything(), anything(), anything()) + ).never(); + }); + + it("prompts to install and runs the install command on accept", async () => { + (window as any).showInformationMessage = async () => + "Install AI tools"; + const manager = createManager(); + + await manager.initialize(); + + assert.ok( + executed.some((e) => e.command === "databricks.aitools.install") + ); + assert.strictEqual( + storedState["databricks.aitools.installPrompted"], + true + ); + }); + + it("does not install when the prompt is dismissed", async () => { + (window as any).showInformationMessage = async () => undefined; + const manager = createManager(); + + await manager.initialize(); + + assert.ok( + !executed.some( + (e) => e.command === "databricks.aitools.install" + ) + ); + assert.strictEqual( + storedState["databricks.aitools.installPrompted"], + true + ); + }); + + it("does not prompt again once prompted", async () => { + storedState["databricks.aitools.installPrompted"] = true; + let prompted = false; + (window as any).showInformationMessage = async () => { + prompted = true; + return undefined; + }; + const manager = createManager(); + + await manager.initialize(); + + assert.strictEqual(prompted, false); + }); + }); +}); diff --git a/packages/databricks-vscode/src/aitools/AiToolsManager.ts b/packages/databricks-vscode/src/aitools/AiToolsManager.ts new file mode 100644 index 000000000..d618a8720 --- /dev/null +++ b/packages/databricks-vscode/src/aitools/AiToolsManager.ts @@ -0,0 +1,510 @@ +import {readFile} from "fs/promises"; +import path from "path"; +import { + commands, + Disposable, + EventEmitter, + ProgressLocation, + window, +} from "vscode"; +import {logging} from "@databricks/sdk-experimental"; +import { + AiToolsListResult, + AiToolsScope, + CliWrapper, + ProcessError, +} from "../cli/CliWrapper"; +import {StateStorage} from "../vscode-objs/StateStorage"; +import {WorkspaceFolderManager} from "../vscode-objs/WorkspaceFolderManager"; +import {Telemetry} from "../telemetry"; +import {AiToolsInstallSource, Events} from "../telemetry/constants"; +import {Loggers} from "../logger"; +import {FileUtils, HostUtils} from "../utils"; + +/** Cursor marketplace numeric ID for the Databricks plugin. */ +const CURSOR_PLUGIN_ID = "26723531"; + +/** Relative path of the aitools state file within an install root. */ +const STATE_FILE_RELATIVE_PATH = path.join( + ".databricks", + "aitools", + "skills", + ".state.json" +); + +/** Where AI tools are installed, or undefined if not installed. */ +export type AiToolsInstallLocation = AiToolsScope | undefined; + +/** The status of the update check. */ +export type AiToolsUpdateStatus = + | "unknown" + | "checking" + | "updating" + | "upToDate" + | "updateAvailable" + | "error"; + +export interface AiToolsState { + installLocation: AiToolsInstallLocation; + updateStatus: AiToolsUpdateStatus; + /** The installed AI tools release version, if known. */ + version?: string; + /** + * True when the last install detection failed with an unexpected error + * (e.g. a permission/IO error reading the state file) rather than the state + * file simply being absent. Distinguishes "genuinely not installed" from + * "couldn't determine install state". + */ + detectError?: boolean; +} + +/** + * Owns all non-UI logic for the Databricks AI tools feature: detecting whether + * tools are installed (and where), running install/update via the CLI, checking + * for available updates, and caching the resolved install location. + */ +export class AiToolsManager implements Disposable { + private disposables: Disposable[] = []; + private readonly onDidChangeEmitter = new EventEmitter(); + public readonly onDidChange = this.onDidChangeEmitter.event; + + private _installLocation: AiToolsInstallLocation; + private _updateStatus: AiToolsUpdateStatus = "unknown"; + private _version: string | undefined; + private _detectError = false; + + constructor( + private readonly cli: CliWrapper, + private readonly stateStorage: StateStorage, + private readonly workspaceFolderManager: WorkspaceFolderManager, + private readonly telemetry: Telemetry + ) { + this._installLocation = this.stateStorage.get( + "databricks.aitools.installLocation" + ); + this.refreshCursorPluginContext(); + this.refreshInstalledContext(); + } + + get state(): AiToolsState { + return { + installLocation: this._installLocation, + updateStatus: this._updateStatus, + version: this._version, + detectError: this._detectError, + }; + } + + get isInstalled(): boolean { + return this._installLocation !== undefined; + } + + /** + * Whether the "add Databricks plugin to Cursor" affordance should be shown: + * only in Cursor, and only if we haven't already prompted for it. + * (Cursor exposes no way to query real plugin state, so this is best-effort.) + */ + get shouldOfferCursorPlugin(): boolean { + return ( + HostUtils.isCursor() && + !this.stateStorage.get("databricks.aitools.cursorPluginPrompted") + ); + } + + /** + * Open Cursor's marketplace install modal for the Databricks plugin, and + * remember that we prompted (to hide the affordance afterwards). We can't + * confirm the user actually added it — only that we opened the modal. + * + * This is decoupled from the CLI install: it opens Cursor's in-app + * marketplace install modal, which is independent of the skills install. + * Any failure here is logged but never propagated, so it can't break the + * install flow when run in parallel. + */ + async addCursorPlugin(): Promise { + try { + await commands.executeCommand( + "workbench.action.openMarketplaceEditor", + { + pluginId: CURSOR_PLUGIN_ID, + openInstallModal: true, + skipTracking: true, + } + ); + } catch (e) { + logging.NamedLogger.getOrCreate(Loggers.Extension).error( + "Failed to open the Cursor marketplace for the Databricks plugin", + e + ); + return; + } + await this.stateStorage.set( + "databricks.aitools.cursorPluginPrompted", + true + ); + this.refreshCursorPluginContext(); + this.onDidChangeEmitter.fire(); + } + + private refreshCursorPluginContext() { + commands.executeCommand( + "setContext", + "databricks.context.aitools.showCursorPlugin", + this.shouldOfferCursorPlugin + ); + } + + /** + * Sync the `databricks.context.aitools.installed` when-context key with the + * current install state, so the command palette can show Install vs. + * Uninstall appropriately. + */ + private refreshInstalledContext() { + commands.executeCommand( + "setContext", + "databricks.context.aitools.installed", + this.isInstalled + ); + } + + dispose() { + this.disposables.forEach((d) => d.dispose()); + this.onDidChangeEmitter.dispose(); + } + + /** + * Whether a workspace folder is open. Project-scope operations need one (the + * skills install into `.databricks/aitools/skills` under the folder); global + * operations run against the home dir and do not. + */ + get hasProjectFolder(): boolean { + // `activeProjectUri` throws when no folder is active; treat that as + // "no project folder" rather than propagating. + try { + return this.workspaceFolderManager.activeProjectUri !== undefined; + } catch { + return false; + } + } + + private get projectRoot(): string { + return this.workspaceFolderManager.activeProjectUri.fsPath; + } + + /** + * Working directory for a CLI invocation, chosen by scope: the project root + * for `project`, the home dir for `global`. Only `project` requires an open + * workspace folder, so global operations work in a folderless window. + */ + private cwdForScope(scope: AiToolsScope): string { + return scope === "project" ? this.projectRoot : FileUtils.getHomedir(); + } + + private stateFilePath(scope: AiToolsScope): string { + return path.join(this.cwdForScope(scope), STATE_FILE_RELATIVE_PATH); + } + + private async stateFileExists(scope: AiToolsScope): Promise { + try { + await readFile(this.stateFilePath(scope)); + return true; + } catch (e: any) { + if (e?.code === "ENOENT") { + return false; + } + throw e; + } + } + + /** + * Determine whether AI tools are installed by checking for + * `.databricks/aitools/skills/.state.json`, first in the project root and + * then in the user's home directory. Caches and persists the resolved + * location and fires {@link onDidChange}. + */ + async detectInstall(): Promise { + let location: AiToolsInstallLocation; + try { + // Project scope only exists when a folder is open; otherwise skip + // straight to checking the global (home dir) install. + if ( + this.hasProjectFolder && + (await this.stateFileExists("project")) + ) { + location = "project"; + } else if (await this.stateFileExists("global")) { + location = "global"; + } + // Detection succeeded (a definitive present/absent answer). + this._detectError = false; + } catch (e) { + // Unexpected error (e.g. EACCES/EIO reading the state file) rather + // than the file being absent. Don't overwrite the last-known-good + // install location with `undefined` — a transient failure must not + // flip an installed toolset to "not installed". Flag the error so + // the UI can surface a reload affordance instead of the install + // prompt. + logging.NamedLogger.getOrCreate(Loggers.Extension).error( + "Failed to detect Databricks AI tools install state", + e + ); + this._detectError = true; + this.refreshInstalledContext(); + this.onDidChangeEmitter.fire(); + return this._installLocation; + } + + this._installLocation = location; + await this.stateStorage.set( + "databricks.aitools.installLocation", + location + ); + if (location === undefined) { + this._updateStatus = "unknown"; + this._version = undefined; + } + this.refreshInstalledContext(); + this.onDidChangeEmitter.fire(); + return location; + } + + /** + * Entry point run once on activation. Detects the install state and then: + * - if installed, checks for updates and, when one is available, applies it + * automatically (updates are silent — no prompt); + * - if not installed, shows a one-time prompt offering to set them up. + * + * Non-blocking failures are swallowed so activation can't be delayed or + * broken by this best-effort flow. + */ + async initialize(): Promise { + try { + const location = await this.detectInstall(); + if (location === undefined) { + await this.maybePromptInstall(); + return; + } + const status = await this.checkForUpdates(); + if (status === "updateAvailable") { + await this.update(); + } + } catch (e) { + logging.NamedLogger.getOrCreate(Loggers.Extension).error( + "Failed to initialize Databricks AI tools", + e + ); + } + } + + /** + * Show a one-time prompt offering to install Databricks AI tools. If the + * user accepts, run the install flow (which, in Cursor, also opens the + * plugin install modal). The prompt is shown at most once per machine. + */ + async maybePromptInstall(): Promise { + if (this.stateStorage.get("databricks.aitools.installPrompted")) { + return; + } + await this.stateStorage.set("databricks.aitools.installPrompted", true); + + const install = "Install AI tools"; + const choice = await window.showInformationMessage( + "Install Databricks AI tools?", + { + modal: true, + detail: "Get Databricks-aware skills and agent plugins in your editor. You can install them later from the Databricks configuration panel.", + }, + install + ); + if (choice !== install) { + return; + } + // Run the install command so the user picks a scope; the install flow + // itself opens the Cursor plugin modal when running in Cursor. Pass the + // "modal" source so telemetry can distinguish first-load prompt installs + // from manual side-pane installs. + await commands.executeCommand("databricks.aitools.install", "modal"); + } + + /** + * Check whether an update is available by comparing each installed skill's + * version against its latest version (via `aitools list --output json`). + * `aitools update --check` only prints text, so `list` is the reliable + * source of truth. + */ + async checkForUpdates(): Promise { + const scope = this._installLocation; + if (scope === undefined) { + this._updateStatus = "unknown"; + this.onDidChangeEmitter.fire(); + return this._updateStatus; + } + + this._updateStatus = "checking"; + this.onDidChangeEmitter.fire(); + + try { + const result = await this.cli.aitoolsList(this.cwdForScope(scope)); + this._version = result.release; + this._updateStatus = this.computeUpdateStatus(result, scope); + } catch (e) { + logging.NamedLogger.getOrCreate(Loggers.Extension).error( + "Failed to check for Databricks AI tools updates", + e + ); + this._updateStatus = "error"; + } + this.onDidChangeEmitter.fire(); + return this._updateStatus; + } + + private computeUpdateStatus( + result: AiToolsListResult, + scope: AiToolsScope + ): AiToolsUpdateStatus { + const installed = result.skills.filter( + (s) => s.installed[scope] !== undefined + ); + const updateAvailable = installed.some( + (s) => s.installed[scope] !== s.latest_version + ); + return updateAvailable ? "updateAvailable" : "upToDate"; + } + + /** + * Install AI tools for the given scope, showing progress. Re-detects the + * install state and refreshes the update status afterwards. + * + * In Cursor, the plugin install is prompted *in parallel* with the CLI + * install — the two are independent, and the plugin modal must not block or + * break the skills install / UI refresh. + */ + async install( + scope: AiToolsScope, + source?: AiToolsInstallSource + ): Promise { + // Kick off the Cursor plugin prompt in parallel (fire-and-forget; it + // swallows its own errors). Not awaited so it can't gate the CLI flow. + if (this.shouldOfferCursorPlugin) { + void this.addCursorPlugin(); + } + + const recordEvent = this.telemetry.start(Events.AITOOLS_INSTALL); + try { + await window.withProgress( + { + location: ProgressLocation.Notification, + title: "Installing Databricks AI tools", + cancellable: true, + }, + (_progress, token) => + this.cli.aitoolsInstall( + scope, + this.cwdForScope(scope), + token + ) + ); + recordEvent({success: true, scope, source}); + } catch (e) { + recordEvent({success: false, scope, source}); + if (e instanceof ProcessError) { + e.showErrorMessage("Failed to install Databricks AI tools."); + } else { + throw e; + } + return; + } + + await this.detectInstall(); + await this.checkForUpdates(); + } + + /** + * Uninstall AI tools for the current install scope, showing progress. + * Re-detects the install state afterwards. + */ + async uninstall(): Promise { + const scope = this._installLocation; + if (scope === undefined) { + return; + } + const recordEvent = this.telemetry.start(Events.AITOOLS_UNINSTALL); + try { + await window.withProgress( + { + location: ProgressLocation.Notification, + title: "Uninstalling Databricks AI tools", + cancellable: true, + }, + (_progress, token) => + this.cli.aitoolsUninstall( + scope, + this.cwdForScope(scope), + token + ) + ); + recordEvent({success: true, scope}); + } catch (e) { + recordEvent({success: false, scope}); + if (e instanceof ProcessError) { + e.showErrorMessage("Failed to uninstall Databricks AI tools."); + } else { + throw e; + } + return; + } + await this.detectInstall(); + + // Clear the Cursor-plugin flag so that if the user reinstalls the tools + // later they're offered the plugin again (uninstalling the skills does + // not remove the Cursor plugin, but re-offering it is harmless and + // matches the "fresh install" expectation). + await this.stateStorage.set( + "databricks.aitools.cursorPluginPrompted", + false + ); + this.refreshCursorPluginContext(); + } + + /** + * Update AI tools for the current install scope, showing progress. + */ + async update(): Promise { + const scope = this._installLocation; + if (scope === undefined) { + return; + } + const recordEvent = this.telemetry.start(Events.AITOOLS_UPDATE); + this._updateStatus = "updating"; + this.onDidChangeEmitter.fire(); + try { + await window.withProgress( + { + location: ProgressLocation.Notification, + title: "Updating Databricks AI tools", + cancellable: true, + }, + (_progress, token) => + this.cli.aitoolsUpdate( + scope, + this.cwdForScope(scope), + token + ) + ); + recordEvent({success: true, scope}); + } catch (e) { + recordEvent({success: false, scope}); + if (e instanceof ProcessError) { + e.showErrorMessage("Failed to update Databricks AI tools."); + } else { + throw e; + } + } finally { + // Always reconcile the cached update status with the actual CLI + // state, even if the update reported an error (it may have + // partially succeeded). This refreshes the row out of the + // "Click to update" state once the tools are up to date. + await this.checkForUpdates(); + } + } +} diff --git a/packages/databricks-vscode/src/bundle/BundleInitWizard.ts b/packages/databricks-vscode/src/bundle/BundleInitWizard.ts index 69985c091..e5a136140 100644 --- a/packages/databricks-vscode/src/bundle/BundleInitWizard.ts +++ b/packages/databricks-vscode/src/bundle/BundleInitWizard.ts @@ -18,13 +18,15 @@ import {Events, Telemetry} from "../telemetry"; import {escapePathArgument} from "../utils/shellUtils"; import {promptToSelectActiveProjectFolder} from "./activeBundleUtils"; import {WorkspaceFolderManager} from "../vscode-objs/WorkspaceFolderManager"; +import {AiToolsManager} from "../aitools/AiToolsManager"; export class BundleInitWizard { private logger = logging.NamedLogger.getOrCreate(Loggers.Extension); constructor( private cli: CliWrapper, - private telemetry: Telemetry + private telemetry: Telemetry, + private aiToolsManager?: AiToolsManager ) {} public async initNewProject( @@ -33,6 +35,8 @@ export class BundleInitWizard { workspaceFolderManager?: WorkspaceFolderManager ) { const recordEvent = this.telemetry.start(Events.BUNDLE_INIT); + // Whether AI tools are already installed when the project is created. + const hasAiTools = this.aiToolsManager?.isInstalled ?? false; try { const authProvider = await this.configureAuthForBundleInit(existingAuthProvider); @@ -40,13 +44,13 @@ export class BundleInitWizard { this.logger.debug( "No valid auth providers, can't proceed with bundle init wizard" ); - recordEvent({success: false}); + recordEvent({success: false, hasAiTools}); return; } const parentFolder = await this.promptForParentFolder(workspaceUri); if (!parentFolder) { this.logger.debug("No parent folder provided"); - recordEvent({success: false}); + recordEvent({success: false, hasAiTools}); return; } await this.bundleInitInTerminal(parentFolder, authProvider); @@ -54,7 +58,7 @@ export class BundleInitWizard { "Finished bundle init wizard, detecting projects to initialize or open" ); const projects = await getSubProjects(parentFolder); - recordEvent({success: projects.length > 0}); + recordEvent({success: projects.length > 0, hasAiTools}); if (projects.length > 0) { this.logger.debug( `Detected ${projects.length} sub projects after the init wizard, prompting to open one` @@ -78,7 +82,7 @@ export class BundleInitWizard { } return parentFolder; } catch (e) { - recordEvent({success: false}); + recordEvent({success: false, hasAiTools}); throw e; } } diff --git a/packages/databricks-vscode/src/bundle/BundleProjectManager.ts b/packages/databricks-vscode/src/bundle/BundleProjectManager.ts index d6bec2ba1..9d74b8230 100644 --- a/packages/databricks-vscode/src/bundle/BundleProjectManager.ts +++ b/packages/databricks-vscode/src/bundle/BundleProjectManager.ts @@ -19,6 +19,7 @@ import {onError} from "../utils/onErrorDecorator"; import {BundleInitWizard} from "./BundleInitWizard"; import {EventReporter, Events, Telemetry} from "../telemetry"; import {WorkspaceFolderManager} from "../vscode-objs/WorkspaceFolderManager"; +import {AiToolsManager} from "../aitools/AiToolsManager"; import {promptToSelectActiveProjectFolder} from "./activeBundleUtils"; export class BundleProjectManager { @@ -50,7 +51,8 @@ export class BundleProjectManager { private configModel: ConfigModel, private bundleFileSet: BundleFileSet, private workspaceFolderManager: WorkspaceFolderManager, - private telemetry: Telemetry + private telemetry: Telemetry, + private aiToolsManager?: AiToolsManager ) { this.disposables.push( this.workspaceFolderManager.onDidChangeActiveProjectFolder( @@ -332,7 +334,11 @@ export class BundleProjectManager { @onError({popup: {prefix: "Failed to initialize new Databricks project"}}) public async initNewProject() { - const bundleInitWizard = new BundleInitWizard(this.cli, this.telemetry); + const bundleInitWizard = new BundleInitWizard( + this.cli, + this.telemetry, + this.aiToolsManager + ); const authProvider = this.connectionManager.databricksWorkspace?.authProvider; const parentFolder = await bundleInitWizard.initNewProject( diff --git a/packages/databricks-vscode/src/cli/CliWrapper.test.ts b/packages/databricks-vscode/src/cli/CliWrapper.test.ts index bd440c1d2..eec5d3188 100644 --- a/packages/databricks-vscode/src/cli/CliWrapper.test.ts +++ b/packages/databricks-vscode/src/cli/CliWrapper.test.ts @@ -4,9 +4,9 @@ import {workspaceConfigs} from "../vscode-objs/WorkspaceConfigs"; import {promisify} from "node:util"; import {execFile as execFileCb} from "node:child_process"; import {withFile} from "tmp-promise"; -import {writeFile, readFile} from "node:fs/promises"; +import {writeFile, readFile, mkdtemp, rm} from "node:fs/promises"; import {when, spy, reset, instance, mock} from "ts-mockito"; -import {CliWrapper, waitForProcess} from "./CliWrapper"; +import {cancellableExecFile, CliWrapper, waitForProcess} from "./CliWrapper"; import path from "node:path"; import os from "node:os"; import crypto from "node:crypto"; @@ -58,6 +58,26 @@ describe(__filename, function () { assert.ok(result.stdout.indexOf("databricks") > 0); }); + it("aitoolsList returns parsed JSON from the bundled CLI", async () => { + const cli = createCliWrapper(); + const tmpDir = await mkdtemp(path.join(os.tmpdir(), "aitools-cli-")); + try { + const result = await cli.aitoolsList(tmpDir); + // The bundled CLI reports the release and the full skill catalog, + // each with a latest_version and an installed map, even when nothing + // is installed in the (empty) temp dir. + assert.ok(typeof result.release === "string"); + assert.ok(Array.isArray(result.skills)); + assert.ok(result.skills.length > 0); + const skill = result.skills[0]; + assert.ok(typeof skill.name === "string"); + assert.ok(typeof skill.latest_version === "string"); + assert.ok(typeof skill.installed === "object"); + } finally { + await rm(tmpDir, {recursive: true, force: true}); + } + }); + it("should resolve the platform-specific CLI binary name", () => { const cli = createCliWrapper(); const originalPlatform = process.platform; @@ -310,6 +330,29 @@ token = dapitest5678 }); }); +describe("cancellableExecFile closeStdin", () => { + // `cat` with no args reads stdin until EOF. Without closeStdin the child's + // stdin pipe stays open forever and the call hangs; closeStdin sends EOF so + // it completes. This mirrors why `aitools update` hung on launch when it + // prompted for confirmation. + it("completes a stdin-reading process when closeStdin is set", async () => { + const {stdout} = await cancellableExecFile("cat", [], {}, undefined, { + closeStdin: true, + }); + assert.strictEqual(stdout, ""); + }); + + it("hangs on a stdin-reading process without closeStdin", async () => { + const raced = await Promise.race([ + cancellableExecFile("cat", []).then(() => "completed"), + new Promise((resolve) => + setTimeout(() => resolve("timed-out"), 500) + ), + ]); + assert.strictEqual(raced, "timed-out"); + }); +}); + describe("waitForProcess", () => { it("should return correctly formatted stdout and stderr", async () => { const process = new ChildProcess(); diff --git a/packages/databricks-vscode/src/cli/CliWrapper.ts b/packages/databricks-vscode/src/cli/CliWrapper.ts index 17ca9170b..8d05ed05a 100644 --- a/packages/databricks-vscode/src/cli/CliWrapper.ts +++ b/packages/databricks-vscode/src/cli/CliWrapper.ts @@ -45,11 +45,24 @@ function getEscapedCommandAndAgrs( return {cmd, args, options}; } +export interface ExecFileOptions { + /** + * Close the child's stdin immediately after spawning. Node's `execFile` + * gives the child an open stdin pipe that never receives EOF, so any CLI + * command that prompts for confirmation (e.g. `aitools update`) blocks + * forever waiting on input. Ending stdin delivers EOF so the prompt + * resolves instead of hanging. Only set this for non-interactive commands + * that we never feed input to. + */ + closeStdin?: boolean; +} + export async function cancellableExecFile( file: string, args: string[], options: Omit = {}, - cancellationToken?: CancellationToken + cancellationToken?: CancellationToken, + execOptions: ExecFileOptions = {} ): Promise<{ stdout: string; stderr: string; @@ -58,10 +71,15 @@ export async function cancellableExecFile( cancellationToken?.onCancellationRequested(() => abortController.abort()); const signal = abortController.signal; - const res = await promisify(execFileCb)(file, args, { + const promise = promisify(execFileCb)(file, args, { ...options, signal, }); + if (execOptions.closeStdin) { + // `promisify(execFile)` exposes the spawned child on `.child`. + (promise as any).child?.stdin?.end(); + } + const res = await promise; return {stdout: res.stdout.toString(), stderr: res.stderr.toString()}; } @@ -69,7 +87,8 @@ export const execFile = async ( file: string, args: string[], options: Omit = {}, - cancellationToken?: CancellationToken + cancellationToken?: CancellationToken, + execOptions: ExecFileOptions = {} ): Promise<{ stdout: string; stderr: string; @@ -84,7 +103,8 @@ export const execFile = async ( cmd, escapedArgs, escapedOptions, - cancellationToken + cancellationToken, + execOptions ); }; @@ -104,6 +124,27 @@ export interface ConfigEntry { } export type SyncType = "full" | "incremental"; + +export type AiToolsScope = "project" | "global"; + +/** A single skill entry from `databricks aitools list --output json`. */ +export interface AiToolsSkill { + name: string; + // eslint-disable-next-line @typescript-eslint/naming-convention + latest_version: string; + experimental: boolean; + /** + * Installed versions keyed by scope. Empty when the skill is not installed. + * e.g. `{ "project": "0.1.0" }` or `{ "global": "0.1.0" }`. + */ + installed: Partial>; +} + +/** Parsed output of `databricks aitools list --output json`. */ +export interface AiToolsListResult { + release: string; + skills: AiToolsSkill[]; +} export class ProcessError extends Error { constructor( message: string, @@ -468,6 +509,125 @@ export class CliWrapper { return stdout; } + private aitoolsEnv(): Record { + return { + ...EnvVarGenerators.getEnvVarsForCli(this.extensionContext), + ...EnvVarGenerators.getProxyEnvVars(), + }; + } + + /** + * Install Databricks AI tools (skills + agent plugins) for the given scope. + * + * `cwd` should be the project root so that `--scope project` installs into + * `.databricks/aitools/skills` under the workspace. The CLI prints + * human-readable text (it ignores `--output json` for this subcommand), so + * success/failure is determined by the exit code (a non-zero exit rejects + * with a {@link ProcessError}). + */ + @withLogContext(Loggers.Extension) + public async aitoolsInstall( + scope: AiToolsScope, + cwd: string, + cancellationToken?: CancellationToken, + @context ctx?: Context + ): Promise { + const args = ["aitools", "install", "--scope", scope]; + try { + await execFile( + this.cliPath, + args, + {cwd, env: this.aitoolsEnv()}, + cancellationToken, + {closeStdin: true} + ); + } catch (e: any) { + ctx?.logger?.error("Failed to install Databricks AI tools", e); + throw new ProcessError(e.message, e.code ?? null); + } + } + + /** + * Update installed Databricks AI tools for the given scope. + */ + @withLogContext(Loggers.Extension) + public async aitoolsUpdate( + scope: AiToolsScope, + cwd: string, + cancellationToken?: CancellationToken, + @context ctx?: Context + ): Promise { + const args = ["aitools", "update", "--scope", scope]; + try { + await execFile( + this.cliPath, + args, + {cwd, env: this.aitoolsEnv()}, + cancellationToken, + {closeStdin: true} + ); + } catch (e: any) { + ctx?.logger?.error("Failed to update Databricks AI tools", e); + throw new ProcessError(e.message, e.code ?? null); + } + } + + /** + * Uninstall Databricks AI tools for the given scope. + */ + @withLogContext(Loggers.Extension) + public async aitoolsUninstall( + scope: AiToolsScope, + cwd: string, + cancellationToken?: CancellationToken, + @context ctx?: Context + ): Promise { + const args = ["aitools", "uninstall", "--scope", scope]; + try { + await execFile( + this.cliPath, + args, + {cwd, env: this.aitoolsEnv()}, + cancellationToken, + {closeStdin: true} + ); + } catch (e: any) { + ctx?.logger?.error("Failed to uninstall Databricks AI tools", e); + throw new ProcessError(e.message, e.code ?? null); + } + } + + /** + * List Databricks AI tools components as structured JSON. + * + * `aitools list` is the only aitools subcommand that emits real JSON + * (`aitools update --check` and `install` print text). We use it both to + * detect whether an update is available (any installed skill whose + * `installed[scope]` differs from `latest_version`) and to read the current + * release. + */ + @withLogContext(Loggers.Extension) + public async aitoolsList( + cwd: string, + @context ctx?: Context + ): Promise { + const args = ["aitools", "list", "--output", "json"]; + let res; + try { + res = await execFile( + this.cliPath, + args, + {cwd, env: this.aitoolsEnv()}, + undefined, + {closeStdin: true} + ); + } catch (e: any) { + ctx?.logger?.error("Failed to list Databricks AI tools", e); + throw new ProcessError(e.message, e.code ?? null); + } + return JSON.parse(res.stdout) as AiToolsListResult; + } + async getBundleCommandEnvVars( authProvider: AuthProvider, configfilePath?: string diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index fd7e799ec..7443f0f3f 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -3,6 +3,7 @@ import { debug, env, ExtensionContext, + ExtensionMode, extensions, window, workspace, @@ -14,6 +15,8 @@ import {ClusterListDataProvider} from "./cluster/ClusterListDataProvider"; import {ClusterModel} from "./cluster/ClusterModel"; import {ClusterCommands} from "./cluster/ClusterCommands"; import {ConfigurationDataProvider} from "./ui/configuration-view/ConfigurationDataProvider"; +import {AiToolsManager} from "./aitools/AiToolsManager"; +import {AiToolsCommands} from "./aitools/AiToolsCommands"; import {RunCommands} from "./run/RunCommands"; import {DatabricksDebugAdapterFactory} from "./run/DatabricksDebugAdapter"; import {DatabricksWorkflowDebugAdapterFactory} from "./run/DatabricksWorkflowDebugAdapter"; @@ -25,6 +28,7 @@ import {logging} from "@databricks/sdk-experimental"; import {workspaceConfigs} from "./vscode-objs/WorkspaceConfigs"; import { FileUtils, + HostUtils, PackageJsonUtils, TerraformUtils, UrlUtils, @@ -38,6 +42,7 @@ import { } from "./workspace-fs"; import {CustomWhenContext} from "./vscode-objs/CustomWhenContext"; import {StateStorage} from "./vscode-objs/StateStorage"; +import {StateResetCommand} from "./vscode-objs/StateResetCommand"; import path from "node:path"; import {FeatureId, FeatureManager} from "./feature-manager/FeatureManager"; import {EnvironmentDependenciesVerifier} from "./language/EnvironmentDependenciesVerifier"; @@ -97,6 +102,7 @@ export async function activate( ): Promise { customWhenContext.setActivated(false); customWhenContext.setDeploymentState("idle"); + customWhenContext.setIsCursor(HostUtils.isCursor()); const stateStorage = new StateStorage(context); const packageMetadata = await PackageJsonUtils.getMetadata(context); @@ -149,6 +155,82 @@ export async function activate( ) ); + const workspaceFolderManager = new WorkspaceFolderManager( + customWhenContext, + stateStorage + ); + + // AI tools: register before the no-folder early return below, since global + // installs/updates work without an open workspace folder. Project scope is + // gated on a folder being open (see AiToolsCommands / AiToolsManager). + const aiToolsManager = new AiToolsManager( + cli, + stateStorage, + workspaceFolderManager, + telemetry + ); + const aiToolsCommands = new AiToolsCommands(aiToolsManager); + context.subscriptions.push( + aiToolsManager, + aiToolsCommands, + telemetry.registerCommand( + "databricks.aitools.install", + aiToolsCommands.installCommand(), + aiToolsCommands + ), + telemetry.registerCommand( + "databricks.aitools.checkForUpdates", + aiToolsCommands.checkForUpdatesCommand(), + aiToolsCommands + ), + telemetry.registerCommand( + "databricks.aitools.reload", + aiToolsCommands.reloadCommand(), + aiToolsCommands + ), + telemetry.registerCommand( + "databricks.aitools.update", + aiToolsCommands.updateCommand(), + aiToolsCommands + ), + telemetry.registerCommand( + "databricks.aitools.uninstall", + aiToolsCommands.uninstallCommand(), + aiToolsCommands + ), + telemetry.registerCommand( + "databricks.aitools.addCursorPlugin", + aiToolsCommands.addCursorPluginCommand(), + aiToolsCommands + ) + ); + // Detect install state on activation and, if installed, auto-apply any + // available update; otherwise prompt the user (once) to set the tools up. + // Non-blocking so it doesn't delay activation. + aiToolsManager.initialize(); + + // Developer-only "Reset state" command (multi-select of persisted state + // keys). Gated to development builds so it isn't exposed to end users; the + // matching `databricks.context.development` context key gates its palette + // entry (see package.json). + const isDevelopment = context.extensionMode === ExtensionMode.Development; + commands.executeCommand( + "setContext", + "databricks.context.development", + isDevelopment + ); + if (isDevelopment) { + const stateResetCommand = new StateResetCommand(stateStorage); + context.subscriptions.push( + stateResetCommand, + telemetry.registerCommand( + "databricks.developer.resetState", + stateResetCommand.resetCommand(), + stateResetCommand + ) + ); + } + if ( workspace.workspaceFolders === undefined || workspace.workspaceFolders?.length === 0 @@ -159,7 +241,8 @@ export async function activate( async () => { const bundleInitWizard = new BundleInitWizard( cli, - telemetry + telemetry, + aiToolsManager ); await bundleInitWizard.initNewProject(); } @@ -168,15 +251,11 @@ export async function activate( // We show a welcome view when there's no workspace folders, prompting users // to either open a new folder or to initialize a new databricks project. // In both cases we expect the workspace to be reloaded and the extension will - // be activated again. + // be activated again. The AI tools setup affordance is added as a + // viewsWelcome entry (see package.json) since it works without a folder. return undefined; } - const workspaceFolderManager = new WorkspaceFolderManager( - customWhenContext, - stateStorage - ); - // Add the databricks binary to the PATH environment variable in terminals context.environmentVariableCollection.clear(); context.environmentVariableCollection.persistent = false; @@ -387,7 +466,8 @@ export async function activate( configModel, bundleFileSet, workspaceFolderManager, - telemetry + telemetry, + aiToolsManager ); context.subscriptions.push( bundleProjectManager, @@ -759,7 +839,8 @@ export async function activate( configModel, cli, featureManager, - workspaceFolderManager + workspaceFolderManager, + aiToolsManager ); const configurationView = window.createTreeView("configurationView", { treeDataProvider: configurationDataProvider, diff --git a/packages/databricks-vscode/src/telemetry/constants.ts b/packages/databricks-vscode/src/telemetry/constants.ts index 9ed8b89f4..f1260a392 100644 --- a/packages/databricks-vscode/src/telemetry/constants.ts +++ b/packages/databricks-vscode/src/telemetry/constants.ts @@ -25,6 +25,9 @@ export enum Events { DBCONNECT_RUN = "dbconnectRun", OPEN_RESOURCE_EXTERNALLY = "openResourceExternally", PYTHON_ENV_SETUP_DETECTED = "python_env.setup.detected", + AITOOLS_INSTALL = "aitoolsInstall", + AITOOLS_UPDATE = "aitoolsUpdate", + AITOOLS_UNINSTALL = "aitoolsUninstall", } /* eslint-enable @typescript-eslint/naming-convention */ @@ -43,6 +46,13 @@ export type BundleRunType = export type WorkflowTaskType = "python" | "notebook" | "unknown"; export type LaunchType = "run" | "debug"; export type ComputeType = "cluster" | "serverless"; +export type AiToolsScope = "project" | "global"; + +/** + * Where an AI tools install was triggered from: the first-load modal prompt or + * the manual affordance in the configuration side pane. + */ +export type AiToolsInstallSource = "modal" | "sidePane"; // Package-manager / interpreter unions are owned by the pure detection module // (the single source of truth) and re-exported here so the event schema and the @@ -166,9 +176,14 @@ export class EventTypes { [Events.BUNDLE_INIT]: EventType< { success: boolean; + hasAiTools?: boolean; } & DurationMeasurement > = { comment: "Initialize a new bundle project", + hasAiTools: { + comment: + "Whether Databricks AI tools are already installed when the project is initialized", + }, }; [Events.BUNDLE_SUB_PROJECTS]: EventType<{ count: number; @@ -226,6 +241,53 @@ export class EventTypes { comment: "The resource type", }, }; + [Events.AITOOLS_INSTALL]: EventType< + { + success: boolean; + scope: AiToolsScope; + source?: AiToolsInstallSource; + } & DurationMeasurement + > = { + comment: "Install Databricks AI tools", + success: { + comment: "true if the install succeeded, false otherwise", + }, + scope: { + comment: "The install scope (project or global)", + }, + source: { + comment: + "Where the install was triggered from: 'modal' (first-load prompt) or 'sidePane' (manual click in the configuration view)", + }, + }; + [Events.AITOOLS_UPDATE]: EventType< + { + success: boolean; + scope: AiToolsScope; + } & DurationMeasurement + > = { + comment: "Update Databricks AI tools", + success: { + comment: "true if the update succeeded, false otherwise", + }, + scope: { + comment: "The update scope (project or global)", + }, + }; + [Events.AITOOLS_UNINSTALL]: EventType< + { + success: boolean; + scope: AiToolsScope; + } & DurationMeasurement + > = { + comment: "Uninstall Databricks AI tools", + success: { + comment: "true if the uninstall succeeded, false otherwise", + }, + scope: { + comment: "The uninstall scope (project or global)", + }, + }; [Events.PYTHON_ENV_SETUP_DETECTED]: EventType<{ managersDetected: PackageManager[]; primaryManager: PrimaryManager; diff --git a/packages/databricks-vscode/src/test/runTest.ts b/packages/databricks-vscode/src/test/runTest.ts index 1aa23a717..102ad08f8 100644 --- a/packages/databricks-vscode/src/test/runTest.ts +++ b/packages/databricks-vscode/src/test/runTest.ts @@ -34,6 +34,9 @@ async function main() { launchArgs: [tmpDir, "--user-data-dir", tmpDir], extensionTestsEnv: { [EXTENSION_DEVELOPMENT]: "true", + ...(process.env.MOCHA_GREP + ? {MOCHA_GREP: process.env.MOCHA_GREP} + : {}), }, }); } catch (err) { diff --git a/packages/databricks-vscode/src/test/suite.ts b/packages/databricks-vscode/src/test/suite.ts index 2d0a05894..e5ba596a0 100644 --- a/packages/databricks-vscode/src/test/suite.ts +++ b/packages/databricks-vscode/src/test/suite.ts @@ -7,6 +7,8 @@ export async function run(): Promise { const mocha = new Mocha({ ui: "bdd", color: true, + // Optional filter to run a subset of tests locally (no-op when unset). + grep: process.env.MOCHA_GREP, }); // Add files to the test suite diff --git a/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.test.ts b/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.test.ts new file mode 100644 index 000000000..b0beed806 --- /dev/null +++ b/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.test.ts @@ -0,0 +1,149 @@ +/* eslint-disable @typescript-eslint/naming-convention */ + +import assert from "assert"; +import {ThemeIcon} from "vscode"; +import { + AiToolsInstallLocation, + AiToolsManager, + AiToolsUpdateStatus, +} from "../../aitools/AiToolsManager"; +import {resolveProviderResult} from "../../test/utils"; +import {AiToolsComponent} from "./AiToolsComponent"; + +function createManager( + installLocation: AiToolsInstallLocation, + updateStatus: AiToolsUpdateStatus, + version?: string, + detectError?: boolean +): AiToolsManager { + return { + state: {installLocation, updateStatus, version, detectError}, + onDidChange: () => ({dispose() {}}), + } as unknown as AiToolsManager; +} + +async function getRoot(manager: AiToolsManager) { + const component = new AiToolsComponent(manager); + const items = await resolveProviderResult(component.getChildren()); + return items ?? []; +} + +describe(__filename, () => { + it("renders a setup prompt when not installed", async () => { + const items = await getRoot(createManager(undefined, "unknown")); + assert.strictEqual(items.length, 1); + const [row] = items; + assert.strictEqual( + row.contextValue, + "databricks.configuration.aitools.notInstalled" + ); + assert.strictEqual(row.command?.command, "databricks.aitools.install"); + }); + + it("renders a retry row when detection failed with no cached location", async () => { + const items = await getRoot( + createManager(undefined, "unknown", undefined, true) + ); + assert.strictEqual(items.length, 1); + const [row] = items; + assert.strictEqual( + row.contextValue, + "databricks.configuration.aitools.error" + ); + assert.ok(String(row.description).includes("Failed to check")); + assert.strictEqual(row.command?.command, "databricks.aitools.reload"); + assert.strictEqual((row.iconPath as ThemeIcon).id, "warning"); + }); + + it("shows the installed row (not the retry row) when a cached location survives an error", async () => { + // detectError is true but a cached location is preserved -> normal row. + const items = await getRoot( + createManager("project", "upToDate", "0.2.9", true) + ); + const [row] = items; + assert.strictEqual(row.label, "AI tools"); + assert.ok(String(row.description).includes("project")); + }); + + it("renders the installed version in the subtext for a project install", async () => { + const items = await getRoot( + createManager("project", "upToDate", "0.2.9") + ); + assert.strictEqual(items.length, 1); + const [row] = items; + assert.strictEqual(row.label, "AI tools"); + assert.strictEqual( + row.contextValue, + "databricks.configuration.aitools.upToDate" + ); + assert.ok(String(row.description).includes("project")); + assert.ok(String(row.description).includes("v0.2.9")); + }); + + it("falls back to 'Up to date' when the version is unknown", async () => { + const items = await getRoot(createManager("project", "upToDate")); + const [row] = items; + assert.ok(String(row.description).includes("Up to date")); + // Stable states use the robot icon. + assert.strictEqual((row.iconPath as ThemeIcon).id, "hubot"); + }); + + it("renders an update-available row without a click command", async () => { + const items = await getRoot(createManager("global", "updateAvailable")); + const [row] = items; + assert.strictEqual( + row.contextValue, + "databricks.configuration.aitools.updateAvailable" + ); + assert.ok(String(row.description).includes("global")); + assert.strictEqual((row.iconPath as ThemeIcon).id, "hubot"); + // Updates apply automatically; the row is not clickable. + assert.strictEqual(row.command, undefined); + }); + + it("renders an updating spinner while auto-updating", async () => { + const items = await getRoot(createManager("project", "updating")); + const [row] = items; + assert.strictEqual( + row.contextValue, + "databricks.configuration.aitools.updating" + ); + assert.ok(String(row.description).includes("Updating")); + assert.strictEqual((row.iconPath as ThemeIcon).id, "sync~spin"); + }); + + it("does not attach a click command to an up-to-date row", async () => { + const items = await getRoot(createManager("project", "upToDate")); + const [row] = items; + assert.strictEqual(row.command, undefined); + }); + + it("renders a checking spinner while checking for updates", async () => { + const items = await getRoot(createManager("project", "checking")); + const [row] = items; + assert.strictEqual( + row.contextValue, + "databricks.configuration.aitools.checking" + ); + assert.strictEqual((row.iconPath as ThemeIcon).id, "sync~spin"); + }); + + it("uses the generic installed context value for unknown status", async () => { + const items = await getRoot(createManager("project", "unknown")); + const [row] = items; + assert.strictEqual( + row.contextValue, + "databricks.configuration.aitools.installed" + ); + }); + + it("returns nothing for a non-root parent", async () => { + const component = new AiToolsComponent( + createManager("project", "upToDate") + ); + const children = await resolveProviderResult( + component.getChildren({label: "AI tools"}) + ); + assert.deepStrictEqual(children, []); + }); +}); diff --git a/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.ts b/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.ts new file mode 100644 index 000000000..001f0e1bb --- /dev/null +++ b/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.ts @@ -0,0 +1,173 @@ +import {ThemeColor, ThemeIcon, TreeItemCollapsibleState} from "vscode"; +import {BaseComponent} from "./BaseComponent"; +import {ConfigurationTreeItem} from "./types"; +import { + AiToolsManager, + AiToolsUpdateStatus, +} from "../../aitools/AiToolsManager"; + +const AITOOLS_COMPONENT_ID = "AITOOLS"; + +/** + * Robot icon used as the AI tools row's identity, in a theme-aware color: + * blue before the tools are installed, green once they are. + */ +function robotIcon(color: "blue" | "green") { + return new ThemeIcon("hubot", new ThemeColor(`charts.${color}`)); +} + +function getContextValue(key: string) { + return `databricks.configuration.aitools.${key}`; +} + +export class AiToolsComponent extends BaseComponent { + constructor(private readonly aiToolsManager: AiToolsManager) { + super(); + this.disposables.push( + this.aiToolsManager.onDidChange(() => { + this.onDidChangeEmitter.fire(); + }) + ); + } + + private getRoot(): ConfigurationTreeItem[] { + const {installLocation, updateStatus, version, detectError} = + this.aiToolsManager.state; + + // Detection failed with an unexpected error and we have no cached + // location to fall back on -> show a reload affordance rather than + // implying the tools simply aren't installed. + if (installLocation === undefined && detectError) { + return [ + { + label: "AI tools", + id: AITOOLS_COMPONENT_ID, + description: + "Failed to check installation · click to retry", + tooltip: + "Failed to check the Databricks AI tools installation. Click to retry.", + contextValue: getContextValue("error"), + iconPath: new ThemeIcon( + "warning", + new ThemeColor("errorForeground") + ), + collapsibleState: TreeItemCollapsibleState.None, + command: { + title: "Retry AI tools detection", + command: "databricks.aitools.reload", + }, + }, + ]; + } + + // Not installed -> prompt to install. + if (installLocation === undefined) { + return [ + { + label: "Install AI tools", + id: AITOOLS_COMPONENT_ID, + contextValue: getContextValue("notInstalled"), + iconPath: robotIcon("blue"), + collapsibleState: TreeItemCollapsibleState.None, + command: { + title: "Install AI tools", + command: "databricks.aitools.install", + }, + }, + ]; + } + + const {icon, description, state} = getTreeItemsForUpdateStatus( + updateStatus, + version + ); + const items: ConfigurationTreeItem[] = [ + { + label: "AI tools", + id: AITOOLS_COMPONENT_ID, + description: `${installLocation}${ + description ? ` · ${description}` : "" + }`, + tooltip: `AI tools installed (${installLocation})`, + contextValue: getContextValue(state), + iconPath: icon, + collapsibleState: TreeItemCollapsibleState.None, + // Updates are applied automatically on activation, so the row + // is not clickable — it only reflects status. + }, + ]; + + // The "add Databricks plugin to Cursor" action is rendered as an inline + // button on this row (see package.json view/item/context menus), gated + // on the databricks.context.aitools.showCursorPlugin context key. + + return items; + } + + public async getChildren( + parent?: ConfigurationTreeItem + ): Promise { + // Single top-level row (plus an optional Cursor plugin row). This + // feature is independent of the cluster connection state. + if (parent !== undefined) { + return []; + } + return this.getRoot(); + } +} + +function getTreeItemsForUpdateStatus( + status: AiToolsUpdateStatus, + version?: string +): { + icon: ThemeIcon; + description?: string; + state: string; +} { + // When we know the installed release, show it (e.g. "v0.2.9") rather than a + // generic "Up to date" label. + const versionLabel = version ? `v${version.replace(/^v/, "")}` : undefined; + switch (status) { + case "upToDate": + return { + // Stable state: the blue robot is the row's resting identity. + icon: robotIcon("green"), + description: versionLabel ?? "Up to date", + state: "upToDate", + }; + case "updateAvailable": + return { + icon: robotIcon("green"), + description: "Update available", + state: "updateAvailable", + }; + case "checking": + // Transient: spinner conveys in-progress work. + return { + icon: new ThemeIcon("sync~spin"), + description: "Checking for updates", + state: "checking", + }; + case "updating": + return { + icon: new ThemeIcon("sync~spin"), + description: "Updating", + state: "updating", + }; + case "error": + return { + icon: new ThemeIcon( + "warning", + new ThemeColor("errorForeground") + ), + description: "Update check failed", + state: "error", + }; + case "unknown": + default: + return { + icon: robotIcon("green"), + state: "installed", + }; + } +} diff --git a/packages/databricks-vscode/src/ui/configuration-view/ConfigurationDataProvider.ts b/packages/databricks-vscode/src/ui/configuration-view/ConfigurationDataProvider.ts index 2fd63c23c..6c6ad62e5 100644 --- a/packages/databricks-vscode/src/ui/configuration-view/ConfigurationDataProvider.ts +++ b/packages/databricks-vscode/src/ui/configuration-view/ConfigurationDataProvider.ts @@ -23,6 +23,8 @@ import {EnvironmentComponent} from "./EnvironmentComponent"; import {WorkspaceFolderComponent} from "./WorkspaceFolderComponent"; import {WorkspaceFolderManager} from "../../vscode-objs/WorkspaceFolderManager"; import {CodeSynchronizer} from "../../sync"; +import {AiToolsComponent} from "./AiToolsComponent"; +import {AiToolsManager} from "../../aitools/AiToolsManager"; /** * Data provider for the cluster tree view @@ -50,10 +52,12 @@ export class ConfigurationDataProvider private readonly configModel: ConfigModel, private readonly cli: CliWrapper, private readonly featureManager: FeatureManager, - private readonly workspaceFolderManager: WorkspaceFolderManager + private readonly workspaceFolderManager: WorkspaceFolderManager, + private readonly aiToolsManager: AiToolsManager ) { this.components = [ new WorkspaceFolderComponent(this.workspaceFolderManager), + new AiToolsComponent(this.aiToolsManager), new BundleTargetComponent(this.configModel), new AuthTypeComponent( this.connectionManager, diff --git a/packages/databricks-vscode/src/utils/hostUtils.ts b/packages/databricks-vscode/src/utils/hostUtils.ts new file mode 100644 index 000000000..5c352a63b --- /dev/null +++ b/packages/databricks-vscode/src/utils/hostUtils.ts @@ -0,0 +1,10 @@ +import {env} from "vscode"; + +/** + * Detect whether the extension is running inside Cursor rather than plain + * VS Code. Cursor reports `env.appName` as "Cursor" (and `env.appHost` as + * "desktop", same as VS Code), so we match on the app name. + */ +export function isCursor(): boolean { + return /cursor/i.test(env.appName); +} diff --git a/packages/databricks-vscode/src/utils/index.ts b/packages/databricks-vscode/src/utils/index.ts index 70f581626..ed6f6af75 100644 --- a/packages/databricks-vscode/src/utils/index.ts +++ b/packages/databricks-vscode/src/utils/index.ts @@ -6,3 +6,4 @@ export * as PackageJsonUtils from "./packageJsonUtils"; export * as EnvVarGenerators from "./envVarGenerators"; export * as DateUtils from "./DateUtils"; export * as TerraformUtils from "./terraformUtils"; +export * as HostUtils from "./hostUtils"; diff --git a/packages/databricks-vscode/src/vscode-objs/CustomWhenContext.ts b/packages/databricks-vscode/src/vscode-objs/CustomWhenContext.ts index e531b9b5c..9eec1133c 100644 --- a/packages/databricks-vscode/src/vscode-objs/CustomWhenContext.ts +++ b/packages/databricks-vscode/src/vscode-objs/CustomWhenContext.ts @@ -116,4 +116,12 @@ export class CustomWhenContext { value ); } + + setIsCursor(value: boolean) { + commands.executeCommand( + "setContext", + "databricks.context.isCursor", + value + ); + } } diff --git a/packages/databricks-vscode/src/vscode-objs/StateResetCommand.ts b/packages/databricks-vscode/src/vscode-objs/StateResetCommand.ts new file mode 100644 index 000000000..715077afd --- /dev/null +++ b/packages/databricks-vscode/src/vscode-objs/StateResetCommand.ts @@ -0,0 +1,57 @@ +import {Disposable, QuickPickItem, commands, window} from "vscode"; +import {StateStorage, StorageKey} from "./StateStorage"; + +interface StateQuickPickItem extends QuickPickItem { + key: StorageKey; +} + +/** + * Developer-only command that lets you reset individual persisted state keys + * (global or workspace) via a multi-select picker. Useful for re-triggering + * one-time flows (e.g. the AI tools install prompt) without wiping the whole + * profile. Registered only in development builds — see extension activation. + */ +export class StateResetCommand implements Disposable { + private disposables: Disposable[] = []; + + constructor(private readonly stateStorage: StateStorage) {} + + dispose() { + this.disposables.forEach((d) => d.dispose()); + } + + resetCommand() { + return async () => { + const items: StateQuickPickItem[] = this.stateStorage.storageKeys + .map(({key, location}) => ({ + key, + label: key, + description: location, + })) + // Stable, readable ordering. + .sort((a, b) => a.label.localeCompare(b.label)); + + const picked = await window.showQuickPick(items, { + title: "Reset Databricks state", + placeHolder: "Select the state keys to reset", + canPickMany: true, + }); + if (picked === undefined || picked.length === 0) { + return; + } + + for (const item of picked) { + await this.stateStorage.reset(item.key); + } + + const reload = "Reload Window"; + const choice = await window.showInformationMessage( + `Reset ${picked.length} state key(s). Reload the window for the change to fully take effect?`, + reload + ); + if (choice === reload) { + await commands.executeCommand("workbench.action.reloadWindow"); + } + }; + } +} diff --git a/packages/databricks-vscode/src/vscode-objs/StateStorage.test.ts b/packages/databricks-vscode/src/vscode-objs/StateStorage.test.ts new file mode 100644 index 000000000..7d749203b --- /dev/null +++ b/packages/databricks-vscode/src/vscode-objs/StateStorage.test.ts @@ -0,0 +1,85 @@ +import assert from "assert"; +import {ExtensionContext} from "vscode"; +import {StateStorage} from "./StateStorage"; + +class InMemoryMemento { + private store = new Map(); + get(key: string, defaultValue?: any) { + return this.store.has(key) ? this.store.get(key) : defaultValue; + } + async update(key: string, value: any) { + if (value === undefined) { + this.store.delete(key); + } else { + this.store.set(key, value); + } + } + keys() { + return [...this.store.keys()]; + } +} + +function createStorage() { + const globalState = new InMemoryMemento(); + const workspaceState = new InMemoryMemento(); + const context = { + globalState, + workspaceState, + } as unknown as ExtensionContext; + return {storage: new StateStorage(context), globalState, workspaceState}; +} + +describe(__filename, () => { + it("enumerates all storage keys with their location", () => { + const {storage} = createStorage(); + const keys = storage.storageKeys; + + const installPrompted = keys.find( + (k) => k.key === "databricks.aitools.installPrompted" + ); + assert.strictEqual(installPrompted?.location, "global"); + + const bundleTarget = keys.find( + (k) => k.key === "databricks.bundle.target" + ); + assert.strictEqual(bundleTarget?.location, "workspace"); + }); + + it("reset clears the stored value so get returns the default", async () => { + const {storage, globalState} = createStorage(); + + await storage.set("databricks.aitools.installPrompted", true); + assert.strictEqual( + storage.get("databricks.aitools.installPrompted"), + true + ); + + await storage.reset("databricks.aitools.installPrompted"); + + // Raw entry removed, and get falls back to the configured default. + assert.strictEqual( + globalState.get("databricks.aitools.installPrompted"), + undefined + ); + assert.strictEqual( + storage.get("databricks.aitools.installPrompted"), + false + ); + }); + + it("reset targets the correct state object by location", async () => { + const {storage, workspaceState} = createStorage(); + + await storage.set("databricks.bundle.target", "dev"); + assert.strictEqual( + workspaceState.get("databricks.bundle.target"), + "dev" + ); + + await storage.reset("databricks.bundle.target"); + assert.strictEqual( + workspaceState.get("databricks.bundle.target"), + undefined + ); + }); +}); diff --git a/packages/databricks-vscode/src/vscode-objs/StateStorage.ts b/packages/databricks-vscode/src/vscode-objs/StateStorage.ts index 04e756258..4b002cd65 100644 --- a/packages/databricks-vscode/src/vscode-objs/StateStorage.ts +++ b/packages/databricks-vscode/src/vscode-objs/StateStorage.ts @@ -76,9 +76,36 @@ const StorageConfigurations = { location: "global", defaultValue: "0.0.0", }), + + // Caches where Databricks AI tools were installed ("project" or "global") + // so update/list commands know which scope and cwd to use. The presence of + // `.databricks/aitools/skills/.state.json` remains the source of truth for + // "are they installed?"; this only caches the resolved location. + "databricks.aitools.installLocation": withType<"project" | "global">()({ + location: "global", + }), + + // Tracks whether we've prompted the user to add the Databricks plugin to + // Cursor. There's no reliable way to detect whether the plugin was actually + // added, so this only records that we opened the install modal (on first + // install in Cursor, or via the "add plugin" action). Used to hide the + // add-plugin affordance once prompted. + "databricks.aitools.cursorPluginPrompted": withType()({ + location: "global", + defaultValue: false, + }), + + // Tracks whether we've already shown the one-time prompt offering to install + // Databricks AI tools. Set once the prompt has been shown so we don't nag on + // every activation. + "databricks.aitools.installPrompted": withType()({ + location: "global", + defaultValue: false, + }), }; -type Keys = keyof typeof StorageConfigurations; +export type StorageKey = keyof typeof StorageConfigurations; +type Keys = StorageKey; type ValueType = (typeof StorageConfigurations)[K]["_type"]; type GetterReturnType> = D extends {getter: infer G} ? G extends (...args: any[]) => any @@ -178,4 +205,26 @@ export class StateStorage { await this.getStateObject(details.location).update(key, value); this.changeEmitters.get(key)?.emitter.fire(); } + + /** All configured storage keys and where each is persisted. */ + get storageKeys(): Array<{key: Keys; location: "global" | "workspace"}> { + return (Object.keys(StorageConfigurations) as Keys[]).map((key) => ({ + key, + location: StorageConfigurations[key].location, + })); + } + + /** + * Remove the persisted value for a key so it reverts to its default. Unlike + * {@link set}, this clears the raw stored entry (rather than writing a + * value), which is what a "reset state" action wants. + */ + @Mutex.synchronise("mutex") + async reset(key: K) { + const details = StorageConfigurations[key] as KeyInfoWithType< + ValueType + >; + await this.getStateObject(details.location).update(key, undefined); + this.changeEmitters.get(key)?.emitter.fire(); + } } diff --git a/packages/databricks-vscode/src/vscode-objs/WorkspaceFolderManager.ts b/packages/databricks-vscode/src/vscode-objs/WorkspaceFolderManager.ts index e7405e8fe..c3a015fdc 100644 --- a/packages/databricks-vscode/src/vscode-objs/WorkspaceFolderManager.ts +++ b/packages/databricks-vscode/src/vscode-objs/WorkspaceFolderManager.ts @@ -75,7 +75,11 @@ export class WorkspaceFolderManager implements Disposable { } private setIsActiveFileInActiveProject() { - if (this.activeProjectUri === undefined) { + // Use the non-throwing private field: the `activeProjectUri` getter + // throws when no folder is open, and this runs from the constructor + // which may execute in a folderless window. + const activeProjectUri = this._activeProjectUri; + if (activeProjectUri === undefined) { this.customWhenContext.setIsActiveFileInActiveWorkspace(false); return; } @@ -83,13 +87,13 @@ export class WorkspaceFolderManager implements Disposable { const isActiveFileInActiveWorkspace = activeEditor !== undefined && activeEditor.document.uri.fsPath.startsWith( - this.activeProjectUri.fsPath + activeProjectUri.fsPath ); const activeNotebookEditor = window.activeNotebookEditor; const isActiveNotebookInActiveWorkspace = activeNotebookEditor !== undefined && activeNotebookEditor.notebook.uri.fsPath.startsWith( - this.activeProjectUri?.fsPath + activeProjectUri.fsPath ); this.customWhenContext.setIsActiveFileInActiveWorkspace( isActiveFileInActiveWorkspace || isActiveNotebookInActiveWorkspace