From 4b8a5593e9d2364f0749953e26398b7a895a8c68 Mon Sep 17 00:00:00 2001 From: Russell Clarey Date: Thu, 16 Jul 2026 14:57:02 +0200 Subject: [PATCH 1/6] aitools integration --- packages/databricks-vscode/package.json | 80 +++ .../src/aitools/AiToolsCommands.ts | 129 ++++ .../src/aitools/AiToolsManager.test.ts | 658 ++++++++++++++++++ .../src/aitools/AiToolsManager.ts | 510 ++++++++++++++ .../src/bundle/BundleInitWizard.ts | 14 +- .../src/bundle/BundleProjectManager.ts | 10 +- .../src/cli/CliWrapper.test.ts | 47 +- .../databricks-vscode/src/cli/CliWrapper.ts | 168 ++++- packages/databricks-vscode/src/extension.ts | 99 ++- .../src/telemetry/constants.ts | 62 ++ .../databricks-vscode/src/test/runTest.ts | 3 + packages/databricks-vscode/src/test/suite.ts | 2 + .../AiToolsComponent.test.ts | 149 ++++ .../ui/configuration-view/AiToolsComponent.ts | 173 +++++ .../ConfigurationDataProvider.ts | 6 +- .../databricks-vscode/src/utils/hostUtils.ts | 10 + packages/databricks-vscode/src/utils/index.ts | 1 + .../src/vscode-objs/CustomWhenContext.ts | 8 + .../src/vscode-objs/StateResetCommand.ts | 57 ++ .../src/vscode-objs/StateStorage.test.ts | 85 +++ .../src/vscode-objs/StateStorage.ts | 51 +- .../src/vscode-objs/WorkspaceFolderManager.ts | 10 +- 22 files changed, 2305 insertions(+), 27 deletions(-) create mode 100644 packages/databricks-vscode/src/aitools/AiToolsCommands.ts create mode 100644 packages/databricks-vscode/src/aitools/AiToolsManager.test.ts create mode 100644 packages/databricks-vscode/src/aitools/AiToolsManager.ts create mode 100644 packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.test.ts create mode 100644 packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.ts create mode 100644 packages/databricks-vscode/src/utils/hostUtils.ts create mode 100644 packages/databricks-vscode/src/vscode-objs/StateResetCommand.ts create mode 100644 packages/databricks-vscode/src/vscode-objs/StateStorage.test.ts 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 From fb06b04b0ec2a2b8478066be6836a2ddc6960d66 Mon Sep 17 00:00:00 2001 From: Russell Clarey Date: Mon, 20 Jul 2026 11:28:29 +0200 Subject: [PATCH 2/6] tweaks --- .../src/aitools/AiToolsCommands.ts | 16 +++++++++------- .../src/aitools/AiToolsManager.ts | 12 +++++++----- packages/databricks-vscode/src/cli/CliWrapper.ts | 16 +++++++++------- packages/databricks-vscode/src/extension.ts | 2 +- .../ui/configuration-view/AiToolsComponent.ts | 7 ++++--- 5 files changed, 30 insertions(+), 23 deletions(-) diff --git a/packages/databricks-vscode/src/aitools/AiToolsCommands.ts b/packages/databricks-vscode/src/aitools/AiToolsCommands.ts index 51f50adc4..cd01dc6f8 100644 --- a/packages/databricks-vscode/src/aitools/AiToolsCommands.ts +++ b/packages/databricks-vscode/src/aitools/AiToolsCommands.ts @@ -29,10 +29,11 @@ export class AiToolsCommands implements Disposable { } /** - * 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. + * Show the scope picker. Project scope is always listed, but when no + * workspace folder is open it shows a "Requires an open folder" hint and + * cannot be selected (the QuickPick API has no true per-item disable, so we + * ignore its selection). Resolves to the chosen scope, or undefined if the + * picker was dismissed. */ private pickScope(): Promise { const hasFolder = this.aiToolsManager.hasProjectFolder; @@ -42,15 +43,16 @@ export class AiToolsCommands implements Disposable { quickPick.items = [ { label: "$(globe) Global", - detail: "Install AI tools for all projects on this machine", + detail: "Available to you across all projects", scope: "global", }, { label: "$(folder) Project", detail: hasFolder - ? "Install AI tools into this project" + ? "Checked into the repo, shared with everyone on the project" : "Open a folder to install AI tools into a project", - // Render as disabled when there's no folder to install into. + // Hint that project scope needs an open folder (selection of + // this item is ignored in onDidAccept when there's no folder). description: hasFolder ? undefined : "Requires an open folder", scope: "project", }, diff --git a/packages/databricks-vscode/src/aitools/AiToolsManager.ts b/packages/databricks-vscode/src/aitools/AiToolsManager.ts index d618a8720..e25af8372 100644 --- a/packages/databricks-vscode/src/aitools/AiToolsManager.ts +++ b/packages/databricks-vscode/src/aitools/AiToolsManager.ts @@ -269,10 +269,11 @@ export class AiToolsManager implements Disposable { } /** - * Entry point run once on activation. Detects the install state and then: + * Entry point run on activation (and by the error-row retry). 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. + * - if not installed, shows a one-time prompt offering to install them. * * Non-blocking failures are swallowed so activation can't be delayed or * broken by this best-effort flow. @@ -312,7 +313,7 @@ export class AiToolsManager implements Disposable { "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.", + detail: "Get skills and plugins so your coding agents work effectively with Databricks. You can also install them later from the Databricks configuration panel.", }, install ); @@ -421,7 +422,8 @@ export class AiToolsManager implements Disposable { /** * Uninstall AI tools for the current install scope, showing progress. - * Re-detects the install state afterwards. + * Re-detects the install state and clears the Cursor-plugin prompt flag + * afterwards (so a later reinstall re-offers the plugin). */ async uninstall(): Promise { const scope = this._installLocation; @@ -503,7 +505,7 @@ export class AiToolsManager implements Disposable { // 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. + // "Update available" state once the tools are up to date. await this.checkForUpdates(); } } diff --git a/packages/databricks-vscode/src/cli/CliWrapper.ts b/packages/databricks-vscode/src/cli/CliWrapper.ts index 8d05ed05a..435638c69 100644 --- a/packages/databricks-vscode/src/cli/CliWrapper.ts +++ b/packages/databricks-vscode/src/cli/CliWrapper.ts @@ -76,8 +76,9 @@ export async function cancellableExecFile( signal, }); if (execOptions.closeStdin) { - // `promisify(execFile)` exposes the spawned child on `.child`. - (promise as any).child?.stdin?.end(); + // `promisify(execFile)` returns a PromiseWithChild that exposes the + // spawned ChildProcess on `.child`. + promise.child.stdin?.end(); } const res = await promise; return {stdout: res.stdout.toString(), stderr: res.stderr.toString()}; @@ -519,11 +520,12 @@ export class CliWrapper { /** * 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}). + * `cwd` selects the install root: the project root for `--scope project` + * (installs into `.databricks/aitools/skills` under the workspace) or the + * home dir for `--scope global` (see AiToolsManager.cwdForScope). 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( diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index 7443f0f3f..9bdcb9f62 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -205,7 +205,7 @@ export async function activate( ) ); // Detect install state on activation and, if installed, auto-apply any - // available update; otherwise prompt the user (once) to set the tools up. + // available update; otherwise prompt the user (once) to install the tools. // Non-blocking so it doesn't delay activation. aiToolsManager.initialize(); diff --git a/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.ts b/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.ts index 001f0e1bb..8d592b4fc 100644 --- a/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.ts +++ b/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.ts @@ -107,8 +107,8 @@ export class AiToolsComponent extends BaseComponent { 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. + // A single top-level AI tools row with no children. This feature is + // independent of the cluster connection state. if (parent !== undefined) { return []; } @@ -130,7 +130,8 @@ function getTreeItemsForUpdateStatus( switch (status) { case "upToDate": return { - // Stable state: the blue robot is the row's resting identity. + // Stable installed state: the green robot is the row's resting + // identity. icon: robotIcon("green"), description: versionLabel ?? "Up to date", state: "upToDate", From e3411f2acbed062d419871be3a6edffeca7875c6 Mon Sep 17 00:00:00 2001 From: Russell Clarey Date: Thu, 30 Jul 2026 12:56:06 +0200 Subject: [PATCH 3/6] list agents in UI pane --- packages/databricks-vscode/package.json | 20 + .../src/aitools/AiToolsCommands.test.ts | 320 +++++++++++++ .../src/aitools/AiToolsCommands.ts | 104 ++++- .../src/aitools/AiToolsManager.test.ts | 440 +++++++++++++++--- .../src/aitools/AiToolsManager.ts | 323 +++++++++---- .../databricks-vscode/src/cli/CliWrapper.ts | 46 +- packages/databricks-vscode/src/extension.ts | 10 + .../src/telemetry/constants.ts | 42 +- .../AiToolsComponent.test.ts | 226 ++++++++- .../ui/configuration-view/AiToolsComponent.ts | 131 +++++- .../src/vscode-objs/StateResetCommand.ts | 1 - .../src/vscode-objs/StateStorage.test.ts | 16 +- .../src/vscode-objs/StateStorage.ts | 20 +- 13 files changed, 1493 insertions(+), 206 deletions(-) create mode 100644 packages/databricks-vscode/src/aitools/AiToolsCommands.test.ts diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index a53467ea2..01ca70f0a 100644 --- a/packages/databricks-vscode/package.json +++ b/packages/databricks-vscode/package.json @@ -127,6 +127,12 @@ "category": "Databricks", "icon": "$(plug)" }, + { + "command": "databricks.aitools.installAgent", + "title": "Install AI tools for this agent", + "category": "Databricks", + "icon": "$(cloud-download)" + }, { "command": "databricks.developer.resetState", "title": "Reset State (Developer)", @@ -476,6 +482,11 @@ "enablement": "databricks.context.activated && !databricks.context.remoteMode", "category": "Databricks" }, + { + "command": "databricks.logs.show", + "title": "Show logs", + "category": "Databricks" + }, { "command": "databricks.utils.copy", "title": "Copy", @@ -1038,6 +1049,11 @@ "when": "view == configurationView && viewItem =~ /^databricks.configuration.aitools.(installed|upToDate|updateAvailable|checking|updating|error)$/", "group": "navigation@9" }, + { + "command": "databricks.aitools.installAgent", + "when": "view == configurationView && viewItem == databricks.configuration.aitools.agent.notInstalled", + "group": "inline@0" + }, { "command": "databricks.bundle.deployAndRunJob", "when": "view == dabsResourceExplorerView && viewItem =~ /^databricks.bundle.resource=jobs.runnable.*$/ && databricks.context.bundle.deploymentState == idle", @@ -1218,6 +1234,10 @@ "command": "databricks.aitools.reload", "when": "false" }, + { + "command": "databricks.aitools.installAgent", + "when": "false" + }, { "command": "databricks.run.runEditorContents", "when": "resourceLangId == python" diff --git a/packages/databricks-vscode/src/aitools/AiToolsCommands.test.ts b/packages/databricks-vscode/src/aitools/AiToolsCommands.test.ts new file mode 100644 index 000000000..cb87e9bb0 --- /dev/null +++ b/packages/databricks-vscode/src/aitools/AiToolsCommands.test.ts @@ -0,0 +1,320 @@ +/* eslint-disable @typescript-eslint/naming-convention */ + +import assert from "assert"; +import {QuickPick, QuickPickItem, window} from "vscode"; +import {anything, capture, instance, mock, verify, when} from "ts-mockito"; +import { + AiToolsAgentStatus, + AiToolsManager, + CURSOR_AGENT_ID, +} from "./AiToolsManager"; +import {AiToolsCommands} from "./AiToolsCommands"; +import {HostUtils} from "../utils"; + +/** + * A minimal, scriptable stand-in for a VS Code QuickPick. `onAccept` decides + * which items are selected (from the items assigned to the pick) and whether the + * pick is accepted or dismissed, then drives the accept/hide callbacks the way + * the real widget would. + */ +class FakeQuickPick { + title?: string; + placeholder?: string; + canSelectMany = false; + items: readonly QuickPickItem[] = []; + selectedItems: readonly QuickPickItem[] = []; + private acceptCbs: Array<() => void> = []; + private hideCbs: Array<() => void> = []; + public disposed = false; + + constructor( + private readonly onAccept: ( + pick: FakeQuickPick + ) => {selected: readonly QuickPickItem[]} | "dismiss" + ) {} + + onDidAccept(cb: () => void) { + this.acceptCbs.push(cb); + return {dispose() {}}; + } + onDidHide(cb: () => void) { + this.hideCbs.push(cb); + return {dispose() {}}; + } + show() { + const result = this.onAccept(this); + if (result === "dismiss") { + // Dismissed without accepting: only hide fires. + this.hideCbs.forEach((cb) => cb()); + return; + } + this.selectedItems = result.selected; + this.acceptCbs.forEach((cb) => cb()); + } + hide() { + this.hideCbs.forEach((cb) => cb()); + } + dispose() { + this.disposed = true; + } +} + +function agent( + id: string, + displayName: string, + detected: boolean +): AiToolsAgentStatus { + return { + id, + displayName, + type: detected ? "plugin" : "skills-only", + detected, + version: detected ? "0.2.10" : undefined, + }; +} + +describe(__filename, () => { + let mockManager: AiToolsManager; + let originalCreateQuickPick: typeof window.createQuickPick; + let originalIsCursor: typeof HostUtils.isCursor; + let quickPicks: FakeQuickPick[]; + // Behavior applied to each createQuickPick call, in order. + let behaviors: Array< + ( + pick: FakeQuickPick + ) => {selected: readonly QuickPickItem[]} | "dismiss" + >; + + function stubIsCursor(value: boolean) { + (HostUtils as any).isCursor = () => value; + } + + beforeEach(() => { + mockManager = mock(AiToolsManager); + when(mockManager.hasProjectFolder).thenReturn(true); + when(mockManager.listAgents(anything())).thenResolve([]); + + quickPicks = []; + behaviors = []; + originalCreateQuickPick = window.createQuickPick; + (window as any).createQuickPick = () => { + const behavior = behaviors.shift() ?? (() => "dismiss" as const); + const pick = new FakeQuickPick(behavior); + quickPicks.push(pick); + return pick as unknown as QuickPick; + }; + + // Default to plain VS Code; Cursor-specific tests opt in. + originalIsCursor = HostUtils.isCursor; + stubIsCursor(false); + }); + + afterEach(() => { + (window as any).createQuickPick = originalCreateQuickPick; + (HostUtils as any).isCursor = originalIsCursor; + }); + + function createCommands() { + return new AiToolsCommands(instance(mockManager)); + } + + /** Select the scope item matching `scope` in the scope picker. */ + function selectScope(scope: "project" | "global") { + return (pick: FakeQuickPick) => ({ + selected: [ + pick.items.find( + (i) => (i as any).scope === scope + ) as QuickPickItem, + ], + }); + } + + it("shows the agent picker after the scope picker and passes the selection to install", async () => { + when(mockManager.listAgents("global")).thenResolve([ + agent("claude-code", "Claude Code", true), + agent("cursor", "Cursor", false), + agent("codex", "Codex CLI", true), + ]); + // Scope picker: choose global. Agent picker: accept as preselected. + behaviors.push(selectScope("global")); + behaviors.push((pick) => ({ + selected: pick.selectedItems, + })); + + await createCommands().installCommand()("sidePane"); + + // The agent picker is the second QuickPick created (after scope). + assert.strictEqual(quickPicks.length, 2); + const agentPick = quickPicks[1]; + assert.strictEqual(agentPick.canSelectMany, true); + assert.deepStrictEqual( + agentPick.items.map((i) => i.label), + ["Claude Code", "Cursor", "Codex CLI"] + ); + + const [scope, source, agents] = capture(mockManager.install).last(); + assert.strictEqual(scope, "global"); + assert.strictEqual(source, "sidePane"); + // Detected agents are preselected and installed by default. + assert.deepStrictEqual(agents, ["claude-code", "codex"]); + }); + + it("preselects only the detected agents", async () => { + when(mockManager.listAgents("global")).thenResolve([ + agent("claude-code", "Claude Code", true), + agent("cursor", "Cursor", false), + agent("codex", "Codex CLI", true), + ]); + behaviors.push(selectScope("global")); + behaviors.push((pick) => ({selected: pick.selectedItems})); + + await createCommands().installCommand()("sidePane"); + + const agentPick = quickPicks[1]; + assert.deepStrictEqual( + agentPick.selectedItems.map((i) => i.label), + ["Claude Code", "Codex CLI"] + ); + // Detected agents carry a "Detected" hint. + const detectedHints = agentPick.items.map((i) => i.description); + assert.deepStrictEqual(detectedHints, [ + "Detected", + undefined, + "Detected", + ]); + }); + + it("starts the Cursor plugin checked and labels it in Cursor", async () => { + stubIsCursor(true); + when(mockManager.listAgents("global")).thenResolve([ + agent("claude-code", "Claude Code", false), + agent(CURSOR_AGENT_ID, "Cursor", false), + ]); + behaviors.push(selectScope("global")); + behaviors.push((pick) => ({selected: pick.selectedItems})); + + await createCommands().installCommand()("sidePane"); + + const agentPick = quickPicks[1]; + // Cursor starts checked even though it isn't "detected"; Claude does not. + assert.deepStrictEqual( + agentPick.selectedItems.map((i) => i.label), + ["Cursor"] + ); + // Cursor is labelled as the plugin rather than a detected skills install. + const hints = agentPick.items.map((i) => i.description); + assert.deepStrictEqual(hints, [undefined, "Databricks plugin"]); + + const [, , agents] = capture(mockManager.install).last(); + assert.deepStrictEqual(agents, [CURSOR_AGENT_ID]); + }); + + it("does not force the Cursor entry checked outside Cursor", async () => { + stubIsCursor(false); + when(mockManager.listAgents("global")).thenResolve([ + agent(CURSOR_AGENT_ID, "Cursor", false), + ]); + behaviors.push(selectScope("global")); + behaviors.push((pick) => ({selected: pick.selectedItems})); + + await createCommands().installCommand()("sidePane"); + + const agentPick = quickPicks[1]; + assert.deepStrictEqual(agentPick.selectedItems, []); + }); + + it("installs the user's edited selection, not just the detected agents", async () => { + when(mockManager.listAgents("global")).thenResolve([ + agent("claude-code", "Claude Code", true), + agent("cursor", "Cursor", false), + ]); + behaviors.push(selectScope("global")); + // User deselects the detected agent and picks the undetected one. + behaviors.push((pick) => ({ + selected: pick.items.filter((i) => (i as any).agentId === "cursor"), + })); + + await createCommands().installCommand()("sidePane"); + + const [, , agents] = capture(mockManager.install).last(); + assert.deepStrictEqual(agents, ["cursor"]); + }); + + it("skips the agent picker and installs with an empty selection when no agents are reported", async () => { + when(mockManager.listAgents("global")).thenResolve([]); + behaviors.push(selectScope("global")); + + await createCommands().installCommand()("sidePane"); + + // Only the scope picker is created; the agent picker is skipped. + assert.strictEqual(quickPicks.length, 1); + const [scope, , agents] = capture(mockManager.install).last(); + assert.strictEqual(scope, "global"); + assert.deepStrictEqual(agents, []); + }); + + it("cancels the install when the agent picker is dismissed", async () => { + when(mockManager.listAgents("global")).thenResolve([ + agent("claude-code", "Claude Code", true), + ]); + behaviors.push(selectScope("global")); + behaviors.push(() => "dismiss"); + + await createCommands().installCommand()("sidePane"); + + verify(mockManager.install(anything(), anything(), anything())).never(); + }); + + it("does not show the agent picker when the scope picker is dismissed", async () => { + behaviors.push(() => "dismiss"); + + await createCommands().installCommand()("sidePane"); + + assert.strictEqual(quickPicks.length, 1); + verify(mockManager.listAgents(anything())).never(); + verify(mockManager.install(anything(), anything(), anything())).never(); + }); + + it("lists agents for the chosen scope", async () => { + when(mockManager.listAgents("project")).thenResolve([ + agent("claude-code", "Claude Code", true), + ]); + behaviors.push(selectScope("project")); + behaviors.push((pick) => ({selected: pick.selectedItems})); + + await createCommands().installCommand()("sidePane"); + + verify(mockManager.listAgents("project")).once(); + }); + + describe("installAgentCommand", () => { + it("installs the agent recovered from the tree node id", async () => { + when(mockManager.installAgent(anything())).thenResolve(); + + await createCommands().installAgentCommand()({ + id: "AITOOLS.agent.codex", + }); + + const [agentId] = capture(mockManager.installAgent).last(); + assert.strictEqual(agentId, "codex"); + }); + + it("ignores a node without an agent id", async () => { + await createCommands().installAgentCommand()({id: "AITOOLS"}); + await createCommands().installAgentCommand()(undefined); + + verify(mockManager.installAgent(anything())).never(); + }); + }); + + describe("addCursorPluginCommand", () => { + it("prompts the plugin with the 'pluginButton' source", async () => { + when(mockManager.addCursorPlugin(anything())).thenResolve(); + + await createCommands().addCursorPluginCommand()(); + + const [source] = capture(mockManager.addCursorPlugin).last(); + assert.strictEqual(source, "pluginButton"); + }); + }); +}); diff --git a/packages/databricks-vscode/src/aitools/AiToolsCommands.ts b/packages/databricks-vscode/src/aitools/AiToolsCommands.ts index cd01dc6f8..6e493a2e8 100644 --- a/packages/databricks-vscode/src/aitools/AiToolsCommands.ts +++ b/packages/databricks-vscode/src/aitools/AiToolsCommands.ts @@ -1,12 +1,21 @@ import {Disposable, QuickPickItem, window} from "vscode"; -import {AiToolsManager} from "./AiToolsManager"; +import { + AiToolsAgentStatus, + AiToolsManager, + CURSOR_AGENT_ID, +} from "./AiToolsManager"; import {AiToolsScope} from "../cli/CliWrapper"; import {AiToolsInstallSource} from "../telemetry/constants"; +import {HostUtils} from "../utils"; interface ScopeQuickPickItem extends QuickPickItem { scope: AiToolsScope; } +interface AgentQuickPickItem extends QuickPickItem { + agentId: string; +} + export class AiToolsCommands implements Disposable { private disposables: Disposable[] = []; @@ -18,13 +27,19 @@ export class AiToolsCommands implements Disposable { 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. + // init modal passes "initModal"); 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); + const agents = await this.pickAgents(scope); + // Dismissing the agent picker cancels the whole install flow. + if (agents === undefined) { + return; + } + await this.aiToolsManager.install(scope, source, agents); }; } @@ -83,9 +98,71 @@ export class AiToolsCommands implements Disposable { }); } + /** + * Show the agent picker for the chosen scope. Lists every agent the CLI + * knows about and allows selecting multiple; agents already present on the + * machine (`detected`) are preselected. Resolves to the selected agent ids, + * or undefined if the picker was dismissed (which cancels the install). + * + * If the CLI reports no agents (e.g. an older CLI, or a list failure), the + * picker is skipped and we resolve to an empty selection so the install + * falls back to the CLI's default (act on every detected agent). + */ + private async pickAgents( + scope: AiToolsScope + ): Promise { + const agents = await this.aiToolsManager.listAgents(scope); + if (agents.length === 0) { + return []; + } + + const inCursor = HostUtils.isCursor(); + const quickPick = window.createQuickPick(); + quickPick.title = "Install Databricks AI tools"; + quickPick.placeholder = "Choose which coding agents to install for"; + quickPick.canSelectMany = true; + const items: AgentQuickPickItem[] = agents.map( + (agent: AiToolsAgentStatus) => { + // In Cursor, the Cursor entry installs the marketplace plugin + // (a superset of the skills); always start it checked and label + // it as the plugin rather than a "Detected" skills install. + const isCursorPlugin = inCursor && agent.id === CURSOR_AGENT_ID; + return { + label: agent.displayName, + description: isCursorPlugin + ? "Databricks plugin" + : agent.detected + ? "Detected" + : undefined, + agentId: agent.id, + picked: isCursorPlugin || agent.detected, + }; + } + ); + quickPick.items = items; + // canSelectMany does not preselect from `picked` alone; set the initial + // selection explicitly so detected agents start checked. + quickPick.selectedItems = items.filter((i) => i.picked); + + return new Promise((resolve) => { + let resolved: string[] | undefined; + this.disposables.push( + quickPick.onDidAccept(() => { + resolved = quickPick.selectedItems.map((i) => i.agentId); + quickPick.hide(); + }), + quickPick.onDidHide(() => { + resolve(resolved); + quickPick.dispose(); + }) + ); + quickPick.show(); + }); + } + checkForUpdatesCommand() { return async () => { - await this.aiToolsManager.checkForUpdates(); + await this.aiToolsManager.resolveInstalled(); }; } @@ -125,7 +202,24 @@ export class AiToolsCommands implements Disposable { addCursorPluginCommand() { return async () => { - await this.aiToolsManager.addCursorPlugin(); + await this.aiToolsManager.addCursorPlugin("pluginButton"); + }; + } + + /** + * Install a single agent from the Agents list's inline "install" button. The + * tree item is passed as the command argument; its id is + * `AITOOLS.agent.`, from which we recover the agent id. + */ + installAgentCommand() { + const prefix = "AITOOLS.agent."; + return async (node?: {id?: string}) => { + if (node?.id === undefined || !node.id.startsWith(prefix)) { + return; + } + await this.aiToolsManager.installAgent( + node.id.slice(prefix.length) + ); }; } } diff --git a/packages/databricks-vscode/src/aitools/AiToolsManager.test.ts b/packages/databricks-vscode/src/aitools/AiToolsManager.test.ts index fc91c04e0..1867f0f60 100644 --- a/packages/databricks-vscode/src/aitools/AiToolsManager.test.ts +++ b/packages/databricks-vscode/src/aitools/AiToolsManager.test.ts @@ -1,16 +1,31 @@ /* eslint-disable @typescript-eslint/naming-convention */ import assert from "assert"; -import {anything, instance, mock, reset, verify, when} from "ts-mockito"; +import { + anything, + capture, + 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 { + AiToolsAgent, + 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"; +import {Events} from "../telemetry/constants"; +import {AiToolsManager, CURSOR_AGENT_ID} from "./AiToolsManager"; +import {HostUtils} from "../utils"; const STATE_FILE_RELATIVE_PATH = path.join( ".databricks", @@ -24,7 +39,8 @@ function listResult( name: string; latest_version: string; installed: Record; - }> + }>, + agents: AiToolsAgent[] = [] ): AiToolsListResult { return { release: "0.2.9", @@ -32,6 +48,7 @@ function listResult( experimental: false, ...s, })), + agents, }; } @@ -45,6 +62,18 @@ describe(__filename, () => { let projectDir: string; let homeDir: string; let originalHome: string | undefined; + let originalIsCursor: typeof HostUtils.isCursor; + // Every telemetry event recorded during a test (via start()'s recorder or a + // direct recordEvent), so assertions can inspect the emitted properties. + let recordedEvents: Array<{event: string; props: any}>; + + function eventsOfType(event: string) { + return recordedEvents.filter((e) => e.event === event); + } + + function stubIsCursor(value: boolean) { + (HostUtils as any).isCursor = () => value; + } async function writeStateFile(root: string) { const dir = path.join(root, path.dirname(STATE_FILE_RELATIVE_PATH)); @@ -75,15 +104,27 @@ describe(__filename, () => { 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. + // Capture recorded events. start() returns a recorder callback that + // records under the event name; recordEvent records directly. + recordedEvents = []; telemetry = { - start: () => () => {}, + start: (event: string) => (props: any) => { + recordedEvents.push({event, props}); + }, + recordEvent: (event: string, props: any) => { + recordedEvents.push({event, props}); + }, } as unknown as Telemetry; + + // Default to plain VS Code; Cursor-specific tests opt in via + // stubIsCursor(true). + originalIsCursor = HostUtils.isCursor; + stubIsCursor(false); }); afterEach(async () => { process.env.HOME = originalHome; + (HostUtils as any).isCursor = originalIsCursor; reset(mockCli); reset(mockWorkspaceFolderManager); await rm(projectDir, {recursive: true, force: true}); @@ -203,8 +244,7 @@ describe(__filename, () => { ); const manager = createManager(); await manager.detectInstall(); - const status = await manager.checkForUpdates(); - assert.strictEqual(status, "upToDate"); + await manager.resolveInstalled(); assert.strictEqual(manager.state.updateStatus, "upToDate"); }); @@ -226,7 +266,8 @@ describe(__filename, () => { ); const manager = createManager(); await manager.detectInstall(); - assert.strictEqual(await manager.checkForUpdates(), "updateAvailable"); + await manager.resolveInstalled(); + assert.strictEqual(manager.state.updateStatus, "updateAvailable"); }); it("ignores non-installed skills when computing update status", async () => { @@ -249,7 +290,8 @@ describe(__filename, () => { ); const manager = createManager(); await manager.detectInstall(); - assert.strictEqual(await manager.checkForUpdates(), "upToDate"); + await manager.resolveInstalled(); + assert.strictEqual(manager.state.updateStatus, "upToDate"); }); it("reports error when the list command fails", async () => { @@ -257,14 +299,15 @@ describe(__filename, () => { when(mockCli.aitoolsList(anything())).thenReject(new Error("boom")); const manager = createManager(); await manager.detectInstall(); - assert.strictEqual(await manager.checkForUpdates(), "error"); + await manager.resolveInstalled(); + assert.strictEqual(manager.state.updateStatus, "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"); + await manager.resolveInstalled(); + assert.strictEqual(manager.state.updateStatus, "unknown"); verify(mockCli.aitoolsList(anything())).never(); }); @@ -330,61 +373,310 @@ describe(__filename, () => { } }); - it("clears the Cursor plugin flag on uninstall so it is re-offered", async () => { - await writeStateFile(projectDir); - 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.aitoolsUninstall("project", anything(), anything()) + mockCli.aitoolsInstall("global", anything(), anything(), anything()) ).thenCall(async () => { - await rm(path.join(projectDir, STATE_FILE_RELATIVE_PATH), { - force: true, - }); + 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.detectInstall(); - await manager.uninstall(); + await manager.install("global"); - assert.strictEqual( - storedState["databricks.aitools.cursorPluginPrompted"], - false - ); + verify( + mockCli.aitoolsInstall("global", anything(), anything(), anything()) + ).once(); + assert.strictEqual(manager.state.installLocation, "global"); + assert.strictEqual(manager.state.updateStatus, "upToDate"); + }); + + describe("Cursor install flow", () => { + let originalExecuteCommand: typeof commands.executeCommand; + let executed: Array<{command: string; args: any[]}>; + + beforeEach(() => { + stubIsCursor(true); + originalExecuteCommand = commands.executeCommand; + executed = []; + (commands as any).executeCommand = async ( + command: string, + ...args: any[] + ) => { + executed.push({command, args}); + }; + }); + + afterEach(() => { + (commands as any).executeCommand = originalExecuteCommand; + }); + + function openedCursorPlugin() { + return executed.some( + (e) => e.command === "workbench.action.openMarketplaceEditor" + ); + } + + it("opens the plugin modal and strips cursor from the CLI --agents", async () => { + await writeStateFile(projectDir); + when( + mockCli.aitoolsInstall( + "project", + anything(), + anything(), + anything() + ) + ).thenResolve(); + when(mockCli.aitoolsList(anything())).thenResolve( + listResult([ + { + name: "databricks-core", + latest_version: "0.1.0", + installed: {project: "0.1.0"}, + }, + ]) + ); + const manager = createManager(); + + await manager.install("project", "sidePane", [ + "claude-code", + CURSOR_AGENT_ID, + ]); + + assert.ok( + openedCursorPlugin(), + "expected the plugin modal to open" + ); + const [, , , cliAgents] = capture(mockCli.aitoolsInstall).last(); + // cursor is never passed to the CLI; the rest are. + assert.deepStrictEqual(cliAgents, ["claude-code"]); + + // The install event records only the CLI agents (cursor stripped) + // and flags the plugin separately. + const [installEvent] = eventsOfType(Events.AITOOLS_INSTALL); + assert.deepStrictEqual(installEvent.props.agents, ["claude-code"]); + assert.strictEqual(installEvent.props.cursorPlugin, true); + // Prompting the plugin is recorded too, inheriting the install's + // source. + const [pluginEvent] = eventsOfType( + Events.AITOOLS_CURSOR_PLUGIN_PROMPT + ); + assert.strictEqual(pluginEvent.props.success, true); + assert.strictEqual(pluginEvent.props.source, "sidePane"); + }); + + it("skips the CLI install when only the Cursor plugin is selected", async () => { + await writeStateFile(projectDir); + const manager = createManager(); + + await manager.install("project", "sidePane", [CURSOR_AGENT_ID]); + + assert.ok( + openedCursorPlugin(), + "expected the plugin modal to open" + ); + // No skills to install via the CLI -> the CLI is not invoked (an + // empty --agents would wrongly act on every detected agent). + verify( + mockCli.aitoolsInstall( + anything(), + anything(), + anything(), + anything() + ) + ).never(); + + // The plugin-only install is still recorded, with no CLI agents. + const [installEvent] = eventsOfType(Events.AITOOLS_INSTALL); + assert.ok(installEvent, "expected an install event"); + assert.strictEqual(installEvent.props.success, true); + assert.deepStrictEqual(installEvent.props.agents, []); + assert.strictEqual(installEvent.props.cursorPlugin, true); + }); + + it("does not open the plugin modal when cursor is not selected", async () => { + await writeStateFile(projectDir); + when( + mockCli.aitoolsInstall( + "project", + anything(), + anything(), + anything() + ) + ).thenResolve(); + when(mockCli.aitoolsList(anything())).thenResolve( + listResult([ + { + name: "databricks-core", + latest_version: "0.1.0", + installed: {project: "0.1.0"}, + }, + ]) + ); + const manager = createManager(); + + await manager.install("project", "sidePane", ["claude-code"]); + + assert.ok( + !openedCursorPlugin(), + "did not expect the plugin modal to open" + ); + const [, , , cliAgents] = capture(mockCli.aitoolsInstall).last(); + assert.deepStrictEqual(cliAgents, ["claude-code"]); + }); + }); + + describe("addCursorPlugin", () => { + let originalExecuteCommand: typeof commands.executeCommand; + + afterEach(() => { + (commands as any).executeCommand = originalExecuteCommand; + }); + + it("records a successful plugin prompt with the given source", async () => { + originalExecuteCommand = commands.executeCommand; + (commands as any).executeCommand = async () => {}; + const manager = createManager(); + + await manager.addCursorPlugin("pluginButton"); + + const [pluginEvent] = eventsOfType( + Events.AITOOLS_CURSOR_PLUGIN_PROMPT + ); + assert.strictEqual(pluginEvent.props.success, true); + assert.strictEqual(pluginEvent.props.source, "pluginButton"); + }); + + it("records a failed plugin prompt when opening the modal throws", async () => { + originalExecuteCommand = commands.executeCommand; + (commands as any).executeCommand = async (command: string) => { + // Only the marketplace modal fails; setContext etc. are no-ops. + if (command === "workbench.action.openMarketplaceEditor") { + throw new Error("no marketplace"); + } + }; + const manager = createManager(); + + await manager.addCursorPlugin(); + + assert.deepStrictEqual( + eventsOfType(Events.AITOOLS_CURSOR_PLUGIN_PROMPT).map( + (e) => e.props.success + ), + [false] + ); + }); }); - it("does not clear the Cursor plugin flag when uninstall fails", async () => { + it("installs a single agent into the current scope and refreshes status", async () => { await writeStateFile(projectDir); - storedState["databricks.aitools.cursorPluginPrompted"] = true; when( - mockCli.aitoolsUninstall("project", anything(), anything()) - ).thenReject(new ProcessError("boom", 1)); + mockCli.aitoolsInstall( + "project", + anything(), + anything(), + anything() + ) + ).thenResolve(); + 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.uninstall(); + await manager.installAgent("codex"); - // The state file still exists (uninstall failed), so the flag must be - // left untouched. - assert.strictEqual( - storedState["databricks.aitools.cursorPluginPrompted"], - true - ); + const [scope, , , agents] = capture(mockCli.aitoolsInstall).last(); + assert.strictEqual(scope, "project"); + assert.deepStrictEqual(agents, ["codex"]); + verify(mockCli.aitoolsList(anything())).once(); + + // The install event records which agent was installed. + const [installEvent] = eventsOfType(Events.AITOOLS_INSTALL); + assert.strictEqual(installEvent.props.success, true); + assert.strictEqual(installEvent.props.source, "sidePane"); + assert.deepStrictEqual(installEvent.props.agents, ["codex"]); }); - it("does not call the CLI when uninstalling with nothing installed", async () => { + it("does not install an agent when nothing is installed", async () => { const manager = createManager(); await manager.detectInstall(); - await manager.uninstall(); + + await manager.installAgent("codex"); + verify( - mockCli.aitoolsUninstall(anything(), anything(), anything()) + mockCli.aitoolsInstall( + anything(), + 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); - } + it("still refreshes the panel when a single-agent install fails", async () => { + // A partial install (e.g. an agent whose CLI is missing) often still + // installed some tools, so the panel must reconcile with the real state + // rather than staying stale. + await writeStateFile(projectDir); + when( + mockCli.aitoolsInstall( + "project", + anything(), + anything(), + anything() + ) + ).thenReject(new ProcessError("boom", 1)); + 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.installAgent("codex"); + + // resolveInstalled ran despite the failure. + verify(mockCli.aitoolsList(anything())).once(); + assert.strictEqual(manager.state.updateStatus, "upToDate"); + }); + + it("still refreshes the panel when the install command fails", async () => { + when( + mockCli.aitoolsInstall("global", anything(), anything(), anything()) + ).thenCall(async () => { + // Simulate a partial install: some tools landed before the failure. + await writeStateFile(homeDir); + throw new ProcessError("boom", 1); + }); when(mockCli.aitoolsList(anything())).thenResolve( listResult([ { @@ -396,9 +688,10 @@ describe(__filename, () => { ); const manager = createManager(); - await manager.install("global"); + await manager.install("global", "sidePane", ["codex"]); - verify(mockCli.aitoolsInstall("global", anything(), anything())).once(); + // detectInstall + resolveInstalled ran despite the failure, so the row + // reflects the tools that actually installed. assert.strictEqual(manager.state.installLocation, "global"); assert.strictEqual(manager.state.updateStatus, "upToDate"); }); @@ -463,10 +756,11 @@ describe(__filename, () => { installed: {project: "0.1.0"}, }, ], + agents: [], }); const manager = createManager(); await manager.detectInstall(); - await manager.checkForUpdates(); + await manager.resolveInstalled(); assert.strictEqual(manager.state.version, "0.3.1"); }); @@ -483,7 +777,7 @@ describe(__filename, () => { ); const manager = createManager(); await manager.detectInstall(); - await manager.checkForUpdates(); + await manager.resolveInstalled(); assert.strictEqual(manager.state.version, "0.2.9"); // Uninstalling / re-detecting with no state file clears the version. @@ -508,7 +802,12 @@ describe(__filename, () => { it("installs global against the home dir without touching projectRoot", async () => { when( - mockCli.aitoolsInstall("global", anything(), anything()) + mockCli.aitoolsInstall( + "global", + anything(), + anything(), + anything() + ) ).thenCall(async () => { await writeStateFile(homeDir); }); @@ -527,7 +826,12 @@ describe(__filename, () => { await manager.install("global"); verify( - mockCli.aitoolsInstall("global", homeDir, anything()) + mockCli.aitoolsInstall( + "global", + homeDir, + anything(), + anything() + ) ).once(); assert.strictEqual(manager.state.installLocation, "global"); }); @@ -618,31 +922,51 @@ describe(__filename, () => { assert.ok( executed.some((e) => e.command === "databricks.aitools.install") ); - assert.strictEqual( - storedState["databricks.aitools.installPrompted"], + // Accepting the install must not set the opt-out flag. + assert.notStrictEqual( + storedState["databricks.aitools.hideInstallPrompt"], true ); }); - it("does not install when the prompt is dismissed", async () => { + it("does not opt out when the prompt is merely dismissed", async () => { (window as any).showInformationMessage = async () => undefined; const manager = createManager(); await manager.initialize(); + assert.ok( + !executed.some( + (e) => e.command === "databricks.aitools.install" + ) + ); + // A plain dismissal leaves the prompt eligible to reappear. + assert.notStrictEqual( + storedState["databricks.aitools.hideInstallPrompt"], + true + ); + }); + + it("opts out permanently when the user picks 'Don't show again'", async () => { + (window as any).showInformationMessage = async () => + "Don't show again"; + const manager = createManager(); + + await manager.initialize(); + assert.ok( !executed.some( (e) => e.command === "databricks.aitools.install" ) ); assert.strictEqual( - storedState["databricks.aitools.installPrompted"], + storedState["databricks.aitools.hideInstallPrompt"], true ); }); - it("does not prompt again once prompted", async () => { - storedState["databricks.aitools.installPrompted"] = true; + it("does not prompt again once opted out", async () => { + storedState["databricks.aitools.hideInstallPrompt"] = true; let prompted = false; (window as any).showInformationMessage = async () => { prompted = true; diff --git a/packages/databricks-vscode/src/aitools/AiToolsManager.ts b/packages/databricks-vscode/src/aitools/AiToolsManager.ts index e25af8372..a6c11f424 100644 --- a/packages/databricks-vscode/src/aitools/AiToolsManager.ts +++ b/packages/databricks-vscode/src/aitools/AiToolsManager.ts @@ -9,21 +9,33 @@ import { } from "vscode"; import {logging} from "@databricks/sdk-experimental"; import { - AiToolsListResult, + AiToolsAgent, AiToolsScope, + AiToolsSkill, 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 { + AiToolsCursorPluginSource, + 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"; +/** + * The Cursor agent id as reported by `aitools list`. In Cursor we install the + * marketplace plugin (a superset of the skills) instead of passing this to the + * CLI, so it's filtered out of any `--agents` selection. + */ +export const CURSOR_AGENT_ID = "cursor"; + /** Relative path of the aitools state file within an install root. */ const STATE_FILE_RELATIVE_PATH = path.join( ".databricks", @@ -44,6 +56,14 @@ export type AiToolsUpdateStatus = | "updateAvailable" | "error"; +export interface AiToolsAgentStatus { + displayName: string; + id: string; + type: "plugin" | "skills-only"; + detected: boolean; + version?: string; +} + export interface AiToolsState { installLocation: AiToolsInstallLocation; updateStatus: AiToolsUpdateStatus; @@ -56,6 +76,30 @@ export interface AiToolsState { * "couldn't determine install state". */ detectError?: boolean; + agents: AiToolsAgentStatus[]; +} + +function computeUpdateStatus( + skills: AiToolsSkill[], + scope: AiToolsScope +): AiToolsUpdateStatus { + const updateAvailable = skills.some( + (s) => s.installed[scope] && s.installed[scope] !== s.latest_version + ); + return updateAvailable ? "updateAvailable" : "upToDate"; +} + +function computeAgentsStatuses( + agents: AiToolsAgent[], + scope: AiToolsScope +): AiToolsAgentStatus[] { + return agents.map((agent) => ({ + displayName: agent.display_name, + id: agent.name, + type: agent.managed ? "plugin" : "skills-only", + detected: agent.detected, + version: agent.installed[scope]?.version, + })); } /** @@ -72,6 +116,7 @@ export class AiToolsManager implements Disposable { private _updateStatus: AiToolsUpdateStatus = "unknown"; private _version: string | undefined; private _detectError = false; + private _agents: AiToolsAgentStatus[] = []; constructor( private readonly cli: CliWrapper, @@ -92,6 +137,7 @@ export class AiToolsManager implements Disposable { updateStatus: this._updateStatus, version: this._version, detectError: this._detectError, + agents: this._agents, }; } @@ -100,28 +146,18 @@ export class AiToolsManager implements Disposable { } /** - * 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. + * Open Cursor's marketplace install modal for the Databricks plugin. 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 { + async addCursorPlugin(source?: AiToolsCursorPluginSource): Promise { + const recordEvent = this.telemetry.start( + Events.AITOOLS_CURSOR_PLUGIN_PROMPT + ); try { await commands.executeCommand( "workbench.action.openMarketplaceEditor", @@ -131,26 +167,26 @@ export class AiToolsManager implements Disposable { skipTracking: true, } ); + recordEvent({success: true, source}); } catch (e) { + recordEvent({success: false, source}); 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(); } + /** + * Whether the "add Databricks plugin to Cursor" button should be visible on + * the top-level AI tools row: always, when running in Cursor, so the user + * can (re-)open the plugin modal at any time. + */ private refreshCursorPluginContext() { commands.executeCommand( "setContext", "databricks.context.aitools.showCursorPlugin", - this.shouldOfferCursorPlugin + HostUtils.isCursor() ); } @@ -285,8 +321,8 @@ export class AiToolsManager implements Disposable { await this.maybePromptInstall(); return; } - const status = await this.checkForUpdates(); - if (status === "updateAvailable") { + await this.resolveInstalled(); + if (this._updateStatus === "updateAvailable") { await this.update(); } } catch (e) { @@ -298,33 +334,70 @@ export class AiToolsManager implements Disposable { } /** - * 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. + * Show the 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 reappears on each activation until the user either installs the + * tools or opts out via "Don't show again" (which sets the + * `hideInstallPrompt` flag). A plain dismissal (Escape/Cancel) does not opt + * the user out, so the offer can resurface later. */ async maybePromptInstall(): Promise { - if (this.stateStorage.get("databricks.aitools.installPrompted")) { + if (this.stateStorage.get("databricks.aitools.hideInstallPrompt")) { return; } - await this.stateStorage.set("databricks.aitools.installPrompted", true); const install = "Install AI tools"; + const dontShowAgain = "Don't show again"; const choice = await window.showInformationMessage( "Install Databricks AI tools?", { modal: true, detail: "Get skills and plugins so your coding agents work effectively with Databricks. You can also install them later from the Databricks configuration panel.", }, - install + install, + dontShowAgain ); + if (choice === dontShowAgain) { + // The user declined and asked not to be prompted again. + await this.stateStorage.set( + "databricks.aitools.hideInstallPrompt", + true + ); + return; + } 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"); + // "initModal" source so telemetry can distinguish first-load prompt + // installs from manual side-pane installs. + await commands.executeCommand( + "databricks.aitools.install", + "initModal" + ); + } + + /** + * List the coding agents known to the CLI for the given scope (via + * `aitools list --output json`), used to populate the install-time agent + * picker. `detected` marks agents already present on the machine so the UI + * can preselect them. Returns an empty array on any failure — the install + * flow then falls back to the CLI's default (act on every detected agent). + */ + async listAgents(scope: AiToolsScope): Promise { + try { + const result = await this.cli.aitoolsList(this.cwdForScope(scope)); + return computeAgentsStatuses(result.agents, scope); + } catch (e) { + logging.NamedLogger.getOrCreate(Loggers.Extension).error( + "Failed to list Databricks AI tools agents", + e + ); + return []; + } } /** @@ -333,12 +406,12 @@ export class AiToolsManager implements Disposable { * `aitools update --check` only prints text, so `list` is the reliable * source of truth. */ - async checkForUpdates(): Promise { + async resolveInstalled(): Promise { const scope = this._installLocation; if (scope === undefined) { this._updateStatus = "unknown"; this.onDidChangeEmitter.fire(); - return this._updateStatus; + return; } this._updateStatus = "checking"; @@ -347,7 +420,8 @@ export class AiToolsManager implements Disposable { try { const result = await this.cli.aitoolsList(this.cwdForScope(scope)); this._version = result.release; - this._updateStatus = this.computeUpdateStatus(result, scope); + this._updateStatus = computeUpdateStatus(result.skills, scope); + this._agents = computeAgentsStatuses(result.agents, scope); } catch (e) { logging.NamedLogger.getOrCreate(Loggers.Extension).error( "Failed to check for Databricks AI tools updates", @@ -356,38 +430,51 @@ export class AiToolsManager implements Disposable { 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. + * In Cursor, selecting the Cursor agent means "install the Databricks + * marketplace plugin" (a superset of the Cursor skills), not "install the + * cursor skills via the CLI". So when running in Cursor and `cursor` is in + * the selection, we open the plugin modal in parallel (fire-and-forget) and + * strip `cursor` from the CLI `--agents` list — we never pass + * `--agents cursor`. */ async install( scope: AiToolsScope, - source?: AiToolsInstallSource + source?: AiToolsInstallSource, + agents?: string[] ): 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(); + let cliAgents = agents; + let cursorPlugin = false; + if (HostUtils.isCursor() && agents !== undefined) { + if (agents.includes(CURSOR_AGENT_ID)) { + // 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. The plugin prompt inherits the + // install's source ('initModal' / 'sidePane'). + cursorPlugin = true; + void this.addCursorPlugin(source); + } + cliAgents = agents.filter((a) => a !== CURSOR_AGENT_ID); + // The user picked *only* the Cursor plugin: there are no skills to + // install via the CLI. Record the install (the plugin) and bail out + // before the CLI call — passing an empty `--agents` list would make + // the CLI act on every detected agent, which is not what was chosen. + if (agents.length > 0 && cliAgents.length === 0) { + this.telemetry.recordEvent(Events.AITOOLS_INSTALL, { + duration: 0, + success: true, + scope, + source, + agents: cliAgents, + cursorPlugin, + }); + return; + } } const recordEvent = this.telemetry.start(Events.AITOOLS_INSTALL); @@ -402,28 +489,99 @@ export class AiToolsManager implements Disposable { this.cli.aitoolsInstall( scope, this.cwdForScope(scope), - token + token, + cliAgents ) ); - recordEvent({success: true, scope, source}); + recordEvent({ + success: true, + scope, + source, + agents: cliAgents, + cursorPlugin, + }); } catch (e) { - recordEvent({success: false, scope, source}); + recordEvent({ + success: false, + scope, + source, + agents: cliAgents, + cursorPlugin, + }); if (e instanceof ProcessError) { - e.showErrorMessage("Failed to install Databricks AI tools."); + e.showErrorMessage( + "Failed to install Databricks AI tools.", + "databricks.logs.show" + ); } else { throw e; } - return; + // Fall through: a failed install often still installed some tools + // (e.g. one agent's CLI was missing), so refresh the panel to + // reflect whatever actually landed rather than leaving it stale. } await this.detectInstall(); - await this.checkForUpdates(); + await this.resolveInstalled(); + } + + /** + * Install a single coding agent into the current install scope, showing + * progress. Used by the per-agent "install" button in the Agents list to + * add an agent that wasn't installed alongside the others. Re-resolves the + * agent statuses afterwards so the row flips to "installed". + */ + async installAgent(agentId: string): Promise { + const scope = this._installLocation; + if (scope === undefined) { + return; + } + const recordEvent = this.telemetry.start(Events.AITOOLS_INSTALL); + try { + await window.withProgress( + { + location: ProgressLocation.Notification, + title: "Installing Databricks AI tools agent", + cancellable: true, + }, + (_progress, token) => + this.cli.aitoolsInstall( + scope, + this.cwdForScope(scope), + token, + [agentId] + ) + ); + recordEvent({ + success: true, + scope, + source: "sidePane", + agents: [agentId], + }); + } catch (e) { + recordEvent({ + success: false, + scope, + source: "sidePane", + agents: [agentId], + }); + if (e instanceof ProcessError) { + e.showErrorMessage( + "Failed to install Databricks AI tools agent.", + "databricks.logs.show" + ); + } else { + throw e; + } + // Fall through to refresh the panel: the install may have partially + // succeeded, so reconcile the row with the real CLI state. + } + await this.resolveInstalled(); } /** * Uninstall AI tools for the current install scope, showing progress. - * Re-detects the install state and clears the Cursor-plugin prompt flag - * afterwards (so a later reinstall re-offers the plugin). + * Re-detects the install state afterwards. */ async uninstall(): Promise { const scope = this._installLocation; @@ -449,23 +607,17 @@ export class AiToolsManager implements Disposable { } catch (e) { recordEvent({success: false, scope}); if (e instanceof ProcessError) { - e.showErrorMessage("Failed to uninstall Databricks AI tools."); + e.showErrorMessage( + "Failed to uninstall Databricks AI tools.", + "databricks.logs.show" + ); } else { throw e; } - return; + // Fall through: a failed uninstall may have removed some tools, so + // re-detect to reflect the real state rather than leaving it stale. } 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(); } /** @@ -497,7 +649,10 @@ export class AiToolsManager implements Disposable { } catch (e) { recordEvent({success: false, scope}); if (e instanceof ProcessError) { - e.showErrorMessage("Failed to update Databricks AI tools."); + e.showErrorMessage( + "Failed to update Databricks AI tools.", + "databricks.logs.show" + ); } else { throw e; } @@ -506,7 +661,7 @@ export class AiToolsManager implements Disposable { // state, even if the update reported an error (it may have // partially succeeded). This refreshes the row out of the // "Update available" state once the tools are up to date. - await this.checkForUpdates(); + await this.resolveInstalled(); } } } diff --git a/packages/databricks-vscode/src/cli/CliWrapper.ts b/packages/databricks-vscode/src/cli/CliWrapper.ts index 435638c69..ce8d5b9aa 100644 --- a/packages/databricks-vscode/src/cli/CliWrapper.ts +++ b/packages/databricks-vscode/src/cli/CliWrapper.ts @@ -19,6 +19,8 @@ import {Context, context} from "@databricks/sdk-experimental/dist/context"; import {Cloud} from "../utils/constants"; import {EnvVarGenerators, FileUtils, UrlUtils} from "../utils"; import {AuthProvider} from "../configuration/auth/AuthProvider"; +import type {AiToolsScope} from "../telemetry/constants"; +export type {AiToolsScope}; import {removeUndefinedKeys} from "../utils/envVarGenerators"; import {quote} from "shell-quote"; import {BundleVariableModel} from "../bundle/models/BundleVariableModel"; @@ -126,8 +128,6 @@ 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; @@ -141,10 +141,27 @@ export interface AiToolsSkill { installed: Partial>; } +export interface AiToolsAgentInstallation { + version: string; + // eslint-disable-next-line @typescript-eslint/naming-convention + native_scope: string; +} + +/** A single agent entry from `databricks aitools list --output json`. */ +export interface AiToolsAgent { + name: string; + // eslint-disable-next-line @typescript-eslint/naming-convention + display_name: string; + managed: boolean; + detected: boolean; + installed: Partial>; +} + /** Parsed output of `databricks aitools list --output json`. */ export interface AiToolsListResult { release: string; skills: AiToolsSkill[]; + agents: AiToolsAgent[]; } export class ProcessError extends Error { constructor( @@ -154,7 +171,21 @@ export class ProcessError extends Error { super(message); } - showErrorMessage(prefix?: string) { + /** + * Show an error toast for this CLI failure with a "Show Logs" button. + * + * `logsCommand` selects which output channel that button opens. It defaults + * to the bundle logs (`databricks.bundle.showLogs`), since most CLI commands + * are bundle operations, but callers whose command logs elsewhere (e.g. the + * AI tools commands, which log to the "Databricks Logs" channel) can pass + * `databricks.logs.show` so "Show Logs" lands on the right channel. + */ + showErrorMessage( + prefix?: string, + logsCommand: + | "databricks.bundle.showLogs" + | "databricks.logs.show" = "databricks.bundle.showLogs" + ) { if (this.message.includes("no value assigned to required variable")) { window .showErrorMessage( @@ -182,7 +213,7 @@ export class ProcessError extends Error { ) .then((choice) => { if (choice === "Show Logs") { - commands.executeCommand("databricks.bundle.showLogs"); + commands.executeCommand(logsCommand); } }); } @@ -532,9 +563,16 @@ export class CliWrapper { scope: AiToolsScope, cwd: string, cancellationToken?: CancellationToken, + agents?: string[], @context ctx?: Context ): Promise { const args = ["aitools", "install", "--scope", scope]; + // When a specific set of agents is chosen, restrict the install to them + // (`--agents a,b`). With no selection the CLI acts on every detected + // agent, which is the desired default. + if (agents && agents.length > 0) { + args.push("--agents", agents.join(",")); + } try { await execFile( this.cliPath, diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index 9bdcb9f62..23c8f0b23 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -142,6 +142,11 @@ export async function activate( "databricks.bundle.showLogs", () => loggerManager.showOutputChannel("Databricks Bundle Logs"), loggerManager + ), + telemetry.registerCommand( + "databricks.logs.show", + () => loggerManager.showOutputChannel("Databricks Logs"), + loggerManager ) ); @@ -202,6 +207,11 @@ export async function activate( "databricks.aitools.addCursorPlugin", aiToolsCommands.addCursorPluginCommand(), aiToolsCommands + ), + telemetry.registerCommand( + "databricks.aitools.installAgent", + aiToolsCommands.installAgentCommand(), + aiToolsCommands ) ); // Detect install state on activation and, if installed, auto-apply any diff --git a/packages/databricks-vscode/src/telemetry/constants.ts b/packages/databricks-vscode/src/telemetry/constants.ts index f1260a392..f1e7617f9 100644 --- a/packages/databricks-vscode/src/telemetry/constants.ts +++ b/packages/databricks-vscode/src/telemetry/constants.ts @@ -28,6 +28,7 @@ export enum Events { AITOOLS_INSTALL = "aitoolsInstall", AITOOLS_UPDATE = "aitoolsUpdate", AITOOLS_UNINSTALL = "aitoolsUninstall", + AITOOLS_CURSOR_PLUGIN_PROMPT = "aitoolsCursorPluginPrompt", } /* eslint-enable @typescript-eslint/naming-convention */ @@ -49,10 +50,18 @@ 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. + * Where an AI tools install was triggered from: the first-load init modal prompt + * or the manual affordance in the configuration side pane. */ -export type AiToolsInstallSource = "modal" | "sidePane"; +export type AiToolsInstallSource = "initModal" | "sidePane"; + +/** + * Where the Cursor plugin prompt was triggered from: as part of an install flow + * ('initModal' / 'sidePane', matching {@link AiToolsInstallSource}), or the + * standalone "add Databricks plugin to Cursor" button on the AI tools row + * ('pluginButton'). + */ +export type AiToolsCursorPluginSource = AiToolsInstallSource | "pluginButton"; // 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 @@ -246,6 +255,8 @@ export class EventTypes { success: boolean; scope: AiToolsScope; source?: AiToolsInstallSource; + agents?: string[]; + cursorPlugin?: boolean; } & DurationMeasurement > = { comment: "Install Databricks AI tools", @@ -257,7 +268,15 @@ export class EventTypes { }, source: { comment: - "Where the install was triggered from: 'modal' (first-load prompt) or 'sidePane' (manual click in the configuration view)", + "Where the install was triggered from: 'initModal' (first-load prompt) or 'sidePane' (manual click in the configuration view)", + }, + agents: { + comment: + 'The coding agents whose skills were installed via the CLI (the closed set of agent ids, e.g. ["claude-code","cursor"]). Excludes the Cursor plugin, which is tracked separately by cursorPlugin. Undefined when no explicit selection was made (the CLI acts on every detected agent).', + }, + cursorPlugin: { + comment: + "In Cursor, whether the Databricks marketplace plugin (a superset of the Cursor skills) was installed as part of this flow, rather than the cursor skills via the CLI", }, }; [Events.AITOOLS_UPDATE]: EventType< @@ -288,6 +307,21 @@ export class EventTypes { comment: "The uninstall scope (project or global)", }, }; + [Events.AITOOLS_CURSOR_PLUGIN_PROMPT]: EventType<{ + success: boolean; + source?: AiToolsCursorPluginSource; + }> = { + comment: + "Prompted the user to install the Databricks plugin from the Cursor marketplace (opened the install modal). We can only observe that we opened the modal, not whether the user actually added the plugin.", + success: { + comment: + "true if the marketplace modal was opened, false if opening it failed", + }, + source: { + comment: + "Where the plugin prompt was triggered from: 'initModal' (first-load install prompt) or 'sidePane' (install triggered from the configuration view), both via the install flow, or 'pluginButton' (the standalone add-plugin button on the AI tools row)", + }, + }; [Events.PYTHON_ENV_SETUP_DETECTED]: EventType<{ managersDetected: PackageManager[]; primaryManager: PrimaryManager; diff --git a/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.test.ts b/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.test.ts index b0beed806..3a4a643ee 100644 --- a/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.test.ts +++ b/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.test.ts @@ -3,6 +3,7 @@ import assert from "assert"; import {ThemeIcon} from "vscode"; import { + AiToolsAgentStatus, AiToolsInstallLocation, AiToolsManager, AiToolsUpdateStatus, @@ -14,10 +15,11 @@ function createManager( installLocation: AiToolsInstallLocation, updateStatus: AiToolsUpdateStatus, version?: string, - detectError?: boolean + detectError?: boolean, + agents: AiToolsAgentStatus[] = [] ): AiToolsManager { return { - state: {installLocation, updateStatus, version, detectError}, + state: {installLocation, updateStatus, version, detectError, agents}, onDidChange: () => ({dispose() {}}), } as unknown as AiToolsManager; } @@ -28,6 +30,29 @@ async function getRoot(manager: AiToolsManager) { return items ?? []; } +async function getChildrenOf( + manager: AiToolsManager, + parent: {label?: string; id?: string} +) { + const component = new AiToolsComponent(manager); + const items = await resolveProviderResult(component.getChildren(parent)); + return items ?? []; +} + +function agent( + id: string, + displayName: string, + version?: string +): AiToolsAgentStatus { + return { + id, + displayName, + type: version !== undefined ? "plugin" : "skills-only", + detected: version !== undefined, + version, + }; +} + describe(__filename, () => { it("renders a setup prompt when not installed", async () => { const items = await getRoot(createManager(undefined, "unknown")); @@ -62,7 +87,7 @@ describe(__filename, () => { ); const [row] = items; assert.strictEqual(row.label, "AI tools"); - assert.ok(String(row.description).includes("project")); + assert.ok(String(row.tooltip).includes("project")); }); it("renders the installed version in the subtext for a project install", async () => { @@ -76,7 +101,7 @@ describe(__filename, () => { row.contextValue, "databricks.configuration.aitools.upToDate" ); - assert.ok(String(row.description).includes("project")); + assert.ok(String(row.tooltip).includes("project")); assert.ok(String(row.description).includes("v0.2.9")); }); @@ -95,7 +120,8 @@ describe(__filename, () => { row.contextValue, "databricks.configuration.aitools.updateAvailable" ); - assert.ok(String(row.description).includes("global")); + assert.ok(String(row.tooltip).includes("global")); + assert.ok(String(row.description).includes("Update available")); assert.strictEqual((row.iconPath as ThemeIcon).id, "hubot"); // Updates apply automatically; the row is not clickable. assert.strictEqual(row.command, undefined); @@ -142,8 +168,196 @@ describe(__filename, () => { createManager("project", "upToDate") ); const children = await resolveProviderResult( - component.getChildren({label: "AI tools"}) + component.getChildren({label: "AI tools", id: "unknown"}) ); assert.deepStrictEqual(children, []); }); + + // The data provider fans getChildren(parent) out to every component and + // flattens the results, so a foreign parent (another component's node being + // expanded) must never make us re-emit the AITOOLS root row. Doing so would + // register a second element with id "AITOOLS" and throw. + it("does not re-emit the root row for a foreign parent when not installed", async () => { + const component = new AiToolsComponent( + createManager(undefined, "unknown") + ); + const children = await resolveProviderResult( + component.getChildren({label: "Some other node", id: "cluster"}) + ); + assert.deepStrictEqual(children, []); + }); + + it("does not re-emit the root row for a foreign parent when detection errored", async () => { + const component = new AiToolsComponent( + createManager(undefined, "unknown", undefined, true) + ); + const children = await resolveProviderResult( + component.getChildren({label: "Some other node", id: "cluster"}) + ); + assert.deepStrictEqual(children, []); + }); + + describe("agents", () => { + it("renders an Agents node summarizing how many are installed", async () => { + const manager = createManager( + "project", + "upToDate", + "0.2.9", + undefined, + [ + agent("claude", "Claude Code", "1.2.0"), + agent("cursor", "Cursor"), + agent("copilot", "GitHub Copilot", "0.5.0"), + ] + ); + const children = await getChildrenOf(manager, { + label: "AI tools", + id: "AITOOLS", + }); + const agentsNode = children.find((c) => c.id === "AITOOLS.agents"); + assert.ok(agentsNode, "expected an Agents node"); + assert.strictEqual(agentsNode.label, "Agents"); + // Only the two agents with a version count as installed. + assert.strictEqual(agentsNode.description, "2 installed"); + }); + + it("reports 0 installed when no agents have a version", async () => { + const manager = createManager( + "project", + "upToDate", + "0.2.9", + undefined, + [agent("cursor", "Cursor")] + ); + const children = await getChildrenOf(manager, { + label: "AI tools", + id: "AITOOLS", + }); + const agentsNode = children.find((c) => c.id === "AITOOLS.agents"); + assert.strictEqual(agentsNode?.description, "0 installed"); + }); + + it("still renders the Agents node when there are no agents", async () => { + const manager = createManager( + "project", + "upToDate", + "0.2.9", + undefined, + [] + ); + const children = await getChildrenOf(manager, { + label: "AI tools", + id: "AITOOLS", + }); + const agentsNode = children.find((c) => c.id === "AITOOLS.agents"); + assert.ok(agentsNode, "expected an Agents node"); + assert.strictEqual(agentsNode.description, "0 installed"); + }); + + it("lists each agent with its version under the Agents node", async () => { + const manager = createManager( + "project", + "upToDate", + "0.2.9", + undefined, + [ + agent("claude", "Claude Code", "1.2.0"), + agent("cursor", "Cursor"), + ] + ); + const rows = await getChildrenOf(manager, { + label: "Agents", + id: "AITOOLS.agents", + }); + assert.strictEqual(rows.length, 2); + + const [claude, cursor] = rows; + assert.strictEqual(claude.label, "Claude Code"); + assert.strictEqual(claude.id, "AITOOLS.agent.claude"); + assert.strictEqual(claude.description, "1.2.0"); + + // Agents without a version render as "Not installed". + assert.strictEqual(cursor.label, "Cursor"); + assert.strictEqual(cursor.id, "AITOOLS.agent.cursor"); + assert.strictEqual(cursor.description, "Not installed"); + }); + + it("returns no agent rows when the agents list is empty", async () => { + const manager = createManager( + "project", + "upToDate", + "0.2.9", + undefined, + [] + ); + const rows = await getChildrenOf(manager, { + label: "Agents", + id: "AITOOLS.agents", + }); + assert.deepStrictEqual(rows, []); + }); + + it("marks installed agents with a green check and no install affordance", async () => { + const manager = createManager( + "project", + "upToDate", + "0.2.9", + undefined, + [agent("claude", "Claude Code", "1.2.0")] + ); + const [row] = await getChildrenOf(manager, { + label: "Agents", + id: "AITOOLS.agents", + }); + assert.strictEqual((row.iconPath as ThemeIcon).id, "check"); + assert.strictEqual( + row.contextValue, + "databricks.configuration.aitools.agent.installed" + ); + // Installed agents are not clickable. + assert.strictEqual(row.command, undefined); + }); + + it("gives uninstalled agents the install context value used by the inline button", async () => { + const manager = createManager( + "project", + "upToDate", + "0.2.9", + undefined, + [agent("codex", "Codex CLI")] + ); + const [row] = await getChildrenOf(manager, { + label: "Agents", + id: "AITOOLS.agents", + }); + assert.strictEqual(row.iconPath, undefined); + assert.strictEqual( + row.contextValue, + "databricks.configuration.aitools.agent.notInstalled" + ); + }); + + it("makes an uninstalled agent row clickable to install it", async () => { + const manager = createManager( + "project", + "upToDate", + "0.2.9", + undefined, + [agent("codex", "Codex CLI")] + ); + const [row] = await getChildrenOf(manager, { + label: "Agents", + id: "AITOOLS.agents", + }); + assert.strictEqual( + row.command?.command, + "databricks.aitools.installAgent" + ); + // The node id is passed through so a click matches the inline + // button, and the command handler can recover the agent id. + assert.deepStrictEqual(row.command?.arguments, [ + {id: "AITOOLS.agent.codex"}, + ]); + }); + }); }); diff --git a/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.ts b/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.ts index 8d592b4fc..f5968470e 100644 --- a/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.ts +++ b/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.ts @@ -4,14 +4,16 @@ import {ConfigurationTreeItem} from "./types"; import { AiToolsManager, AiToolsUpdateStatus, + CURSOR_AGENT_ID, } from "../../aitools/AiToolsManager"; +import {HostUtils} from "../../utils"; -const AITOOLS_COMPONENT_ID = "AITOOLS"; +const TREE_ICON_ID = "AITOOLS"; + +function getTreeIconId(key: string) { + return `${TREE_ICON_ID}.${key}`; +} -/** - * 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}`)); } @@ -20,6 +22,16 @@ function getContextValue(key: string) { return `databricks.configuration.aitools.${key}`; } +/** + * Agents that should not be listed under the Agents node. In Cursor the Cursor + * plugin is managed via the marketplace modal (surfaced as a button on the + * top-level AI tools row), so listing "Cursor" as an agent here would be + * redundant and confusing. + */ +function isHiddenAgent(agentId: string): boolean { + return agentId === CURSOR_AGENT_ID && HostUtils.isCursor(); +} + export class AiToolsComponent extends BaseComponent { constructor(private readonly aiToolsManager: AiToolsManager) { super(); @@ -41,7 +53,7 @@ export class AiToolsComponent extends BaseComponent { return [ { label: "AI tools", - id: AITOOLS_COMPONENT_ID, + id: TREE_ICON_ID, description: "Failed to check installation · click to retry", tooltip: @@ -60,12 +72,11 @@ export class AiToolsComponent extends BaseComponent { ]; } - // Not installed -> prompt to install. if (installLocation === undefined) { return [ { label: "Install AI tools", - id: AITOOLS_COMPONENT_ID, + id: TREE_ICON_ID, contextValue: getContextValue("notInstalled"), iconPath: robotIcon("blue"), collapsibleState: TreeItemCollapsibleState.None, @@ -84,16 +95,12 @@ export class AiToolsComponent extends BaseComponent { const items: ConfigurationTreeItem[] = [ { label: "AI tools", - id: AITOOLS_COMPONENT_ID, - description: `${installLocation}${ - description ? ` · ${description}` : "" - }`, + id: TREE_ICON_ID, + 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. + collapsibleState: TreeItemCollapsibleState.Collapsed, }, ]; @@ -107,12 +114,95 @@ export class AiToolsComponent extends BaseComponent { public async getChildren( parent?: ConfigurationTreeItem ): Promise { - // A single top-level AI tools row with no children. This feature is - // independent of the cluster connection state. - if (parent !== undefined) { + const {installLocation, version, detectError, agents} = + this.aiToolsManager.state; + // Only the tree root gets the AI tools row. Guarding solely on + // `parent === undefined` is important: ConfigurationDataProvider fans + // every getChildren(parent) call out to all components and flattens the + // results, so returning the root row for a foreign parent (e.g. when a + // cluster/auth node is expanded) would register a second element with + // id "AITOOLS" and throw "Element with id AITOOLS is already + // registered". + if (parent === undefined) { + return this.getRoot(); + } + + // The child rows below only exist under an installed, expandable AI + // tools row. When nothing is installed (or detection errored) the root + // row is non-collapsible, so it is never expanded and these branches + // are unreachable for our own nodes; bail out for any other parent. + if (installLocation === undefined || detectError) { return []; } - return this.getRoot(); + + if (parent.id === TREE_ICON_ID) { + // In Cursor, the Cursor plugin is installed via the marketplace + // modal (the button on the top-level AI tools row), not the CLI, so + // it never appears as a manageable agent under the Agents node. + const visibleAgents = agents.filter((a) => !isHiddenAgent(a.id)); + const installedAgents = visibleAgents.filter( + (a) => a.version !== undefined + ).length; + return [ + { + label: "Scope", + id: getTreeIconId("scope"), + description: installLocation, + collapsibleState: TreeItemCollapsibleState.None, + }, + version !== undefined && { + label: "Version", + id: getTreeIconId("version"), + description: version, + collapsibleState: TreeItemCollapsibleState.None, + }, + agents !== undefined && { + label: "Agents", + id: getTreeIconId("agents"), + description: `${installedAgents} installed`, + collapsibleState: TreeItemCollapsibleState.Expanded, + }, + ].filter(Boolean) as ConfigurationTreeItem[]; + } + if (parent.id === getTreeIconId("agents")) { + return agents + .filter((agent) => !isHiddenAgent(agent.id)) + .map((agent) => { + const id = getTreeIconId(`agent.${agent.id}`); + const installed = agent.version !== undefined; + return { + label: agent.displayName, + id, + description: agent.version ?? "Not installed", + // A green check marks installed agents; uninstalled ones + // get an inline install button (see package.json + // view/item/context) keyed on this context value, and + // clicking the row installs the agent too. + contextValue: getContextValue( + installed ? "agent.installed" : "agent.notInstalled" + ), + iconPath: installed + ? new ThemeIcon( + "check", + new ThemeColor("charts.green") + ) + : undefined, + command: installed + ? undefined + : { + title: "Install AI tools for this agent", + command: "databricks.aitools.installAgent", + // The handler recovers the agent id from the + // node id; pass the node explicitly so a click + // and the inline button behave identically. + arguments: [{id}], + }, + collapsibleState: TreeItemCollapsibleState.None, + }; + }); + } + + return []; } } @@ -130,8 +220,6 @@ function getTreeItemsForUpdateStatus( switch (status) { case "upToDate": return { - // Stable installed state: the green robot is the row's resting - // identity. icon: robotIcon("green"), description: versionLabel ?? "Up to date", state: "upToDate", @@ -143,7 +231,6 @@ function getTreeItemsForUpdateStatus( state: "updateAvailable", }; case "checking": - // Transient: spinner conveys in-progress work. return { icon: new ThemeIcon("sync~spin"), description: "Checking for updates", diff --git a/packages/databricks-vscode/src/vscode-objs/StateResetCommand.ts b/packages/databricks-vscode/src/vscode-objs/StateResetCommand.ts index 715077afd..bebfba345 100644 --- a/packages/databricks-vscode/src/vscode-objs/StateResetCommand.ts +++ b/packages/databricks-vscode/src/vscode-objs/StateResetCommand.ts @@ -28,7 +28,6 @@ export class StateResetCommand implements Disposable { label: key, description: location, })) - // Stable, readable ordering. .sort((a, b) => a.label.localeCompare(b.label)); const picked = await window.showQuickPick(items, { diff --git a/packages/databricks-vscode/src/vscode-objs/StateStorage.test.ts b/packages/databricks-vscode/src/vscode-objs/StateStorage.test.ts index 7d749203b..42388a82f 100644 --- a/packages/databricks-vscode/src/vscode-objs/StateStorage.test.ts +++ b/packages/databricks-vscode/src/vscode-objs/StateStorage.test.ts @@ -34,10 +34,10 @@ describe(__filename, () => { const {storage} = createStorage(); const keys = storage.storageKeys; - const installPrompted = keys.find( - (k) => k.key === "databricks.aitools.installPrompted" + const hideInstallPrompt = keys.find( + (k) => k.key === "databricks.aitools.hideInstallPrompt" ); - assert.strictEqual(installPrompted?.location, "global"); + assert.strictEqual(hideInstallPrompt?.location, "global"); const bundleTarget = keys.find( (k) => k.key === "databricks.bundle.target" @@ -48,21 +48,21 @@ describe(__filename, () => { it("reset clears the stored value so get returns the default", async () => { const {storage, globalState} = createStorage(); - await storage.set("databricks.aitools.installPrompted", true); + await storage.set("databricks.aitools.hideInstallPrompt", true); assert.strictEqual( - storage.get("databricks.aitools.installPrompted"), + storage.get("databricks.aitools.hideInstallPrompt"), true ); - await storage.reset("databricks.aitools.installPrompted"); + await storage.reset("databricks.aitools.hideInstallPrompt"); // Raw entry removed, and get falls back to the configured default. assert.strictEqual( - globalState.get("databricks.aitools.installPrompted"), + globalState.get("databricks.aitools.hideInstallPrompt"), undefined ); assert.strictEqual( - storage.get("databricks.aitools.installPrompted"), + storage.get("databricks.aitools.hideInstallPrompt"), false ); }); diff --git a/packages/databricks-vscode/src/vscode-objs/StateStorage.ts b/packages/databricks-vscode/src/vscode-objs/StateStorage.ts index 4b002cd65..195cd8b81 100644 --- a/packages/databricks-vscode/src/vscode-objs/StateStorage.ts +++ b/packages/databricks-vscode/src/vscode-objs/StateStorage.ts @@ -85,20 +85,12 @@ const StorageConfigurations = { 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()({ + // Tracks whether the user has opted out of the prompt offering to install + // Databricks AI tools. Only set when the user declines the install *and* + // asks not to be shown it again (the "Don't show again" affordance); a plain + // dismissal leaves this false so the prompt can reappear on a later + // activation. + "databricks.aitools.hideInstallPrompt": withType()({ location: "global", defaultValue: false, }), From cda54dbe6fbbd43b7f2a76030bab7ae9119c65d8 Mon Sep 17 00:00:00 2001 From: Russell Clarey Date: Thu, 30 Jul 2026 15:55:56 +0200 Subject: [PATCH 4/6] align with code conventions & improve tests --- .../src/aitools/AiToolsCommands.test.ts | 140 ++++++++- .../src/aitools/AiToolsCommands.ts | 149 ++++++++-- .../src/aitools/AiToolsManager.test.ts | 131 +++------ .../src/aitools/AiToolsManager.ts | 266 +++++++----------- .../src/cli/CliWrapper.test.ts | 80 +++++- packages/databricks-vscode/src/extension.ts | 2 +- .../AiToolsComponent.test.ts | 57 ++++ .../src/utils/hostUtils.test.ts | 39 +++ .../src/vscode-objs/StateResetCommand.test.ts | 130 +++++++++ .../WorkspaceFolderManager.test.ts | 58 +++- 10 files changed, 759 insertions(+), 293 deletions(-) create mode 100644 packages/databricks-vscode/src/utils/hostUtils.test.ts create mode 100644 packages/databricks-vscode/src/vscode-objs/StateResetCommand.test.ts diff --git a/packages/databricks-vscode/src/aitools/AiToolsCommands.test.ts b/packages/databricks-vscode/src/aitools/AiToolsCommands.test.ts index cb87e9bb0..263c786d5 100644 --- a/packages/databricks-vscode/src/aitools/AiToolsCommands.test.ts +++ b/packages/databricks-vscode/src/aitools/AiToolsCommands.test.ts @@ -1,8 +1,15 @@ /* eslint-disable @typescript-eslint/naming-convention */ import assert from "assert"; -import {QuickPick, QuickPickItem, window} from "vscode"; +import { + CancellationToken, + Progress, + QuickPick, + QuickPickItem, + window, +} from "vscode"; import {anything, capture, instance, mock, verify, when} from "ts-mockito"; +import {ProcessError} from "../cli/CliWrapper"; import { AiToolsAgentStatus, AiToolsManager, @@ -76,6 +83,7 @@ function agent( describe(__filename, () => { let mockManager: AiToolsManager; let originalCreateQuickPick: typeof window.createQuickPick; + let originalWithProgress: typeof window.withProgress; let originalIsCursor: typeof HostUtils.isCursor; let quickPicks: FakeQuickPick[]; // Behavior applied to each createQuickPick call, in order. @@ -84,6 +92,12 @@ describe(__filename, () => { pick: FakeQuickPick ) => {selected: readonly QuickPickItem[]} | "dismiss" >; + // A cancellation token whose `isCancellationRequested` is false; enough to + // stand in for the token withProgress would hand the task. + const fakeToken = { + isCancellationRequested: false, + onCancellationRequested: () => ({dispose() {}}), + } as unknown as CancellationToken; function stubIsCursor(value: boolean) { (HostUtils as any).isCursor = () => value; @@ -104,6 +118,17 @@ describe(__filename, () => { return pick as unknown as QuickPick; }; + // Run the progress task synchronously with a non-cancelled token, + // skipping the real notification UI. + originalWithProgress = window.withProgress; + (window as any).withProgress = ( + _options: unknown, + task: ( + progress: Progress, + token: CancellationToken + ) => Thenable + ) => task({report() {}}, fakeToken); + // Default to plain VS Code; Cursor-specific tests opt in. originalIsCursor = HostUtils.isCursor; stubIsCursor(false); @@ -111,6 +136,7 @@ describe(__filename, () => { afterEach(() => { (window as any).createQuickPick = originalCreateQuickPick; + (window as any).withProgress = originalWithProgress; (HostUtils as any).isCursor = originalIsCursor; }); @@ -289,7 +315,9 @@ describe(__filename, () => { describe("installAgentCommand", () => { it("installs the agent recovered from the tree node id", async () => { - when(mockManager.installAgent(anything())).thenResolve(); + when( + mockManager.installAgent(anything(), anything()) + ).thenResolve(); await createCommands().installAgentCommand()({ id: "AITOOLS.agent.codex", @@ -303,7 +331,7 @@ describe(__filename, () => { await createCommands().installAgentCommand()({id: "AITOOLS"}); await createCommands().installAgentCommand()(undefined); - verify(mockManager.installAgent(anything())).never(); + verify(mockManager.installAgent(anything(), anything())).never(); }); }); @@ -317,4 +345,110 @@ describe(__filename, () => { assert.strictEqual(source, "pluginButton"); }); }); + + describe("initializeCommand", () => { + it("shows the install prompt and installs on accept when not installed", async () => { + when(mockManager.initialize()).thenResolve("promptInstall"); + when(mockManager.listAgents("global")).thenResolve([]); + const originalShowInfo = window.showInformationMessage; + (window as any).showInformationMessage = async () => + "Install AI tools"; + behaviors.push(selectScope("global")); + try { + await createCommands().initializeCommand()(); + } finally { + (window as any).showInformationMessage = originalShowInfo; + } + + // Accepting the prompt runs the install flow with the "initModal" + // source (never the opt-out). + const [, source] = capture(mockManager.install).last(); + assert.strictEqual(source, "initModal"); + verify(mockManager.optOutOfInstallPrompt()).never(); + }); + + it("opts out when the prompt is declined with 'Don't show again'", async () => { + when(mockManager.initialize()).thenResolve("promptInstall"); + when(mockManager.optOutOfInstallPrompt()).thenResolve(); + const originalShowInfo = window.showInformationMessage; + (window as any).showInformationMessage = async () => + "Don't show again"; + try { + await createCommands().initializeCommand()(); + } finally { + (window as any).showInformationMessage = originalShowInfo; + } + + verify(mockManager.optOutOfInstallPrompt()).once(); + verify( + mockManager.install(anything(), anything(), anything()) + ).never(); + }); + + it("does not opt out or install on a plain dismissal", async () => { + when(mockManager.initialize()).thenResolve("promptInstall"); + const originalShowInfo = window.showInformationMessage; + (window as any).showInformationMessage = async () => undefined; + try { + await createCommands().initializeCommand()(); + } finally { + (window as any).showInformationMessage = originalShowInfo; + } + + verify(mockManager.optOutOfInstallPrompt()).never(); + verify( + mockManager.install(anything(), anything(), anything()) + ).never(); + }); + + it("applies the update when one is available", async () => { + when(mockManager.initialize()).thenResolve("update"); + when(mockManager.update(anything())).thenResolve(); + + await createCommands().initializeCommand()(); + + verify(mockManager.update(anything())).once(); + }); + + it("does nothing when the action is 'none'", async () => { + when(mockManager.initialize()).thenResolve("none"); + + await createCommands().initializeCommand()(); + + verify(mockManager.update(anything())).never(); + verify( + mockManager.install(anything(), anything(), anything()) + ).never(); + }); + }); + + describe("error handling", () => { + it("surfaces a ProcessError from update as a toast (does not rethrow)", async () => { + const err = new ProcessError("boom", 1); + let shownPrefix: string | undefined; + err.showErrorMessage = (prefix?: string) => { + shownPrefix = prefix; + }; + when(mockManager.update(anything())).thenReject(err); + + // A ProcessError is caught and rendered, not propagated. + await createCommands().updateCommand()(); + + assert.strictEqual( + shownPrefix, + "Failed to update Databricks AI tools." + ); + }); + + it("rethrows a non-ProcessError from update", async () => { + when(mockManager.update(anything())).thenReject( + new Error("unexpected") + ); + + await assert.rejects( + () => createCommands().updateCommand()(), + /unexpected/ + ); + }); + }); }); diff --git a/packages/databricks-vscode/src/aitools/AiToolsCommands.ts b/packages/databricks-vscode/src/aitools/AiToolsCommands.ts index 6e493a2e8..15297aa04 100644 --- a/packages/databricks-vscode/src/aitools/AiToolsCommands.ts +++ b/packages/databricks-vscode/src/aitools/AiToolsCommands.ts @@ -1,10 +1,16 @@ -import {Disposable, QuickPickItem, window} from "vscode"; +import { + CancellationToken, + Disposable, + ProgressLocation, + QuickPickItem, + window, +} from "vscode"; import { AiToolsAgentStatus, AiToolsManager, CURSOR_AGENT_ID, } from "./AiToolsManager"; -import {AiToolsScope} from "../cli/CliWrapper"; +import {AiToolsScope, ProcessError} from "../cli/CliWrapper"; import {AiToolsInstallSource} from "../telemetry/constants"; import {HostUtils} from "../utils"; @@ -25,24 +31,112 @@ export class AiToolsCommands implements Disposable { this.disposables.forEach((d) => d.dispose()); } + /** + * Run a cancellable operation inside a progress notification, surfacing any + * {@link ProcessError} as an error toast whose "Show Logs" button opens the + * "Databricks Logs" channel. Non-`ProcessError` failures propagate. This is + * the single place the AI tools UI wraps the manager's (logic-only) + * install/update/uninstall methods with VS Code chrome. + */ + private async withProgress( + title: string, + errorPrefix: string, + run: (token: CancellationToken) => Promise + ): Promise { + try { + await window.withProgress( + { + location: ProgressLocation.Notification, + title, + cancellable: true, + }, + (_progress, token) => run(token) + ); + } catch (e) { + if (e instanceof ProcessError) { + e.showErrorMessage(errorPrefix, "databricks.logs.show"); + return; + } + throw e; + } + } + + /** + * Run the one-time activation flow: ask the manager what's needed, then + * render it. If not installed (and not opted out), show the install prompt; + * if an update is available, apply it silently with progress. Invoked on + * activation and by the error-row retry. + */ + initializeCommand() { + return async () => { + const action = await this.aiToolsManager.initialize(); + if (action === "promptInstall") { + await this.promptInstall(); + } else if (action === "update") { + await this.runUpdate(); + } + }; + } + + /** + * Show the one-time prompt offering to install Databricks AI tools. On + * accept, run the install flow (passing the "initModal" source so telemetry + * can distinguish first-load prompt installs from manual side-pane ones). On + * "Don't show again", opt the user out; a plain dismissal leaves the offer + * eligible to reappear on a later activation. + */ + private async promptInstall(): Promise { + const install = "Install AI tools"; + const dontShowAgain = "Don't show again"; + const choice = await window.showInformationMessage( + "Install Databricks AI tools?", + { + modal: true, + detail: "Get skills and plugins so your coding agents work effectively with Databricks. You can also install them later from the Databricks configuration panel.", + }, + install, + dontShowAgain + ); + if (choice === dontShowAgain) { + await this.aiToolsManager.optOutOfInstallPrompt(); + return; + } + if (choice !== install) { + return; + } + await this.runInstall("initModal"); + } + installCommand() { // The command may be invoked with a source argument (e.g. the first-load // init modal passes "initModal"); default to "sidePane" for the manual // affordance. return async (source: AiToolsInstallSource = "sidePane") => { - const scope = await this.pickScope(); - if (scope === undefined) { - return; - } - const agents = await this.pickAgents(scope); - // Dismissing the agent picker cancels the whole install flow. - if (agents === undefined) { - return; - } - await this.aiToolsManager.install(scope, source, agents); + await this.runInstall(source); }; } + /** + * Drive the install flow: pick a scope, pick the agents, then run the + * install with a progress notification. Dismissing either picker cancels. + */ + private async runInstall(source: AiToolsInstallSource): Promise { + const scope = await this.pickScope(); + if (scope === undefined) { + return; + } + const agents = await this.pickAgents(scope); + // Dismissing the agent picker cancels the whole install flow. + if (agents === undefined) { + return; + } + await this.withProgress( + "Installing Databricks AI tools", + "Failed to install Databricks AI tools.", + (token) => this.aiToolsManager.install(scope, source, agents, token) + ); + } + /** * Show the scope picker. Project scope is always listed, but when no * workspace folder is open it shows a "Requires an open folder" hint and @@ -167,18 +261,24 @@ export class AiToolsCommands implements Disposable { } /** - * Re-run install detection (and update check if installed). Used by the - * error row to recover after a transient detection failure. + * Re-run the activation flow. Used by the error row to recover after a + * transient detection failure. */ reloadCommand() { - return async () => { - await this.aiToolsManager.initialize(); - }; + return this.initializeCommand(); + } + + private async runUpdate(): Promise { + await this.withProgress( + "Updating Databricks AI tools", + "Failed to update Databricks AI tools.", + (token) => this.aiToolsManager.update(token) + ); } updateCommand() { return async () => { - await this.aiToolsManager.update(); + await this.runUpdate(); }; } @@ -196,7 +296,11 @@ export class AiToolsCommands implements Disposable { if (confirm !== "Uninstall") { return; } - await this.aiToolsManager.uninstall(); + await this.withProgress( + "Uninstalling Databricks AI tools", + "Failed to uninstall Databricks AI tools.", + (token) => this.aiToolsManager.uninstall(token) + ); }; } @@ -217,8 +321,11 @@ export class AiToolsCommands implements Disposable { if (node?.id === undefined || !node.id.startsWith(prefix)) { return; } - await this.aiToolsManager.installAgent( - node.id.slice(prefix.length) + const agentId = node.id.slice(prefix.length); + await this.withProgress( + "Installing Databricks AI tools agent", + "Failed to install Databricks AI tools agent.", + (token) => this.aiToolsManager.installAgent(agentId, token) ); }; } diff --git a/packages/databricks-vscode/src/aitools/AiToolsManager.test.ts b/packages/databricks-vscode/src/aitools/AiToolsManager.test.ts index 1867f0f60..57f81920c 100644 --- a/packages/databricks-vscode/src/aitools/AiToolsManager.test.ts +++ b/packages/databricks-vscode/src/aitools/AiToolsManager.test.ts @@ -10,7 +10,7 @@ import { verify, when, } from "ts-mockito"; -import {commands, Uri, window} from "vscode"; +import {commands, Uri} from "vscode"; import {mkdtemp, mkdir, writeFile, rm} from "fs/promises"; import path from "path"; import os from "os"; @@ -640,7 +640,8 @@ describe(__filename, () => { it("still refreshes the panel when a single-agent install fails", async () => { // A partial install (e.g. an agent whose CLI is missing) often still // installed some tools, so the panel must reconcile with the real state - // rather than staying stale. + // rather than staying stale. The error is rethrown for the command layer + // to surface. await writeStateFile(projectDir); when( mockCli.aitoolsInstall( @@ -662,7 +663,7 @@ describe(__filename, () => { const manager = createManager(); await manager.detectInstall(); - await manager.installAgent("codex"); + await assert.rejects(() => manager.installAgent("codex"), ProcessError); // resolveInstalled ran despite the failure. verify(mockCli.aitoolsList(anything())).once(); @@ -688,7 +689,10 @@ describe(__filename, () => { ); const manager = createManager(); - await manager.install("global", "sidePane", ["codex"]); + await assert.rejects( + () => manager.install("global", "sidePane", ["codex"]), + ProcessError + ); // detectInstall + resolveInstalled ran despite the failure, so the row // reflects the tools that actually installed. @@ -737,7 +741,7 @@ describe(__filename, () => { const manager = createManager(); await manager.detectInstall(); - await manager.update(); + await assert.rejects(() => manager.update(), ProcessError); // The finally block reconciles state even though the update errored. assert.strictEqual(manager.state.updateStatus, "updateAvailable"); @@ -844,55 +848,31 @@ describe(__filename, () => { }); 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 () => { + it("resolves the update status but reports 'update' when installed and behind", 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([ + when(mockCli.aitoolsList(anything())).thenResolve( + listResult([ { name: "databricks-core", latest_version: "0.1.0", - installed: {project: call === 1 ? "0.0.1" : "0.1.0"}, + installed: {project: "0.0.1"}, }, - ]); - }); + ]) + ); const manager = createManager(); - await manager.initialize(); + const action = await manager.initialize(); + // initialize resolves status but leaves applying the update to the + // caller (AiToolsCommands). + assert.strictEqual(action, "update"); + assert.strictEqual(manager.state.updateStatus, "updateAvailable"); verify( - mockCli.aitoolsUpdate("project", anything(), anything()) - ).once(); - assert.strictEqual(manager.state.updateStatus, "upToDate"); + mockCli.aitoolsUpdate(anything(), anything(), anything()) + ).never(); }); - it("does not update when already up to date", async () => { + it("reports 'none' when installed and up to date", async () => { await writeStateFile(projectDir); when(mockCli.aitoolsList(anything())).thenResolve( listResult([ @@ -905,78 +885,33 @@ describe(__filename, () => { ); const manager = createManager(); - await manager.initialize(); - - verify( - mockCli.aitoolsUpdate(anything(), anything(), anything()) - ).never(); + assert.strictEqual(await manager.initialize(), "none"); }); - it("prompts to install and runs the install command on accept", async () => { - (window as any).showInformationMessage = async () => - "Install AI tools"; + it("reports 'promptInstall' when not installed", async () => { const manager = createManager(); - await manager.initialize(); - - assert.ok( - executed.some((e) => e.command === "databricks.aitools.install") - ); - // Accepting the install must not set the opt-out flag. - assert.notStrictEqual( - storedState["databricks.aitools.hideInstallPrompt"], - true - ); + assert.strictEqual(await manager.initialize(), "promptInstall"); }); - it("does not opt out when the prompt is merely dismissed", async () => { - (window as any).showInformationMessage = async () => undefined; + it("reports 'none' when not installed but opted out", async () => { + storedState["databricks.aitools.hideInstallPrompt"] = true; const manager = createManager(); - await manager.initialize(); - - assert.ok( - !executed.some( - (e) => e.command === "databricks.aitools.install" - ) - ); - // A plain dismissal leaves the prompt eligible to reappear. - assert.notStrictEqual( - storedState["databricks.aitools.hideInstallPrompt"], - true - ); + assert.strictEqual(await manager.initialize(), "none"); }); - it("opts out permanently when the user picks 'Don't show again'", async () => { - (window as any).showInformationMessage = async () => - "Don't show again"; + it("records the opt-out via optOutOfInstallPrompt", async () => { const manager = createManager(); + assert.strictEqual(manager.shouldPromptInstall, true); - await manager.initialize(); + await manager.optOutOfInstallPrompt(); - assert.ok( - !executed.some( - (e) => e.command === "databricks.aitools.install" - ) - ); assert.strictEqual( storedState["databricks.aitools.hideInstallPrompt"], true ); - }); - - it("does not prompt again once opted out", async () => { - storedState["databricks.aitools.hideInstallPrompt"] = true; - let prompted = false; - (window as any).showInformationMessage = async () => { - prompted = true; - return undefined; - }; - const manager = createManager(); - - await manager.initialize(); - - assert.strictEqual(prompted, false); + assert.strictEqual(manager.shouldPromptInstall, false); }); }); }); diff --git a/packages/databricks-vscode/src/aitools/AiToolsManager.ts b/packages/databricks-vscode/src/aitools/AiToolsManager.ts index a6c11f424..5ba8ee1fc 100644 --- a/packages/databricks-vscode/src/aitools/AiToolsManager.ts +++ b/packages/databricks-vscode/src/aitools/AiToolsManager.ts @@ -1,19 +1,12 @@ import {readFile} from "fs/promises"; import path from "path"; -import { - commands, - Disposable, - EventEmitter, - ProgressLocation, - window, -} from "vscode"; +import {CancellationToken, commands, Disposable, EventEmitter} from "vscode"; import {logging} from "@databricks/sdk-experimental"; import { AiToolsAgent, AiToolsScope, AiToolsSkill, CliWrapper, - ProcessError, } from "../cli/CliWrapper"; import {StateStorage} from "../vscode-objs/StateStorage"; import {WorkspaceFolderManager} from "../vscode-objs/WorkspaceFolderManager"; @@ -56,6 +49,18 @@ export type AiToolsUpdateStatus = | "updateAvailable" | "error"; +/** + * What activation should do after {@link AiToolsManager.initialize} has detected + * the install state and (if installed) resolved the update status. The manager + * decides *what* is needed but leaves the UI (progress, prompt) to + * {@link AiToolsCommands}: + * - `"promptInstall"` — not installed: offer the one-time install prompt. + * - `"update"` — installed and an update is available: apply it (silently). + * - `"none"` — installed and up to date (or opted out / detection errored): + * nothing to do. + */ +export type AiToolsInitAction = "promptInstall" | "update" | "none"; + export interface AiToolsAgentStatus { displayName: string; id: string; @@ -307,76 +312,52 @@ export class AiToolsManager implements Disposable { /** * Entry point run on activation (and by the error-row retry). 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 install them. + * - if installed, checks for updates and reports whether one is available so + * the caller can apply it (updates are silent — no prompt); + * - if not installed (and the user hasn't opted out), reports that the + * one-time install prompt should be shown. * - * Non-blocking failures are swallowed so activation can't be delayed or - * broken by this best-effort flow. + * Returns the {@link AiToolsInitAction} the caller should take. All UI + * (progress, prompt) is left to {@link AiToolsCommands}. Non-blocking + * failures are swallowed (resolving to `"none"`) so activation can't be + * delayed or broken by this best-effort flow. */ - async initialize(): Promise { + async initialize(): Promise { try { const location = await this.detectInstall(); if (location === undefined) { - await this.maybePromptInstall(); - return; + return this.shouldPromptInstall ? "promptInstall" : "none"; } await this.resolveInstalled(); - if (this._updateStatus === "updateAvailable") { - await this.update(); - } + return this._updateStatus === "updateAvailable" ? "update" : "none"; } catch (e) { logging.NamedLogger.getOrCreate(Loggers.Extension).error( "Failed to initialize Databricks AI tools", e ); + return "none"; } } /** - * Show the 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 reappears on each activation until the user either installs the - * tools or opts out via "Don't show again" (which sets the - * `hideInstallPrompt` flag). A plain dismissal (Escape/Cancel) does not opt - * the user out, so the offer can resurface later. + * Whether the one-time "install AI tools" prompt should be offered: true + * unless the user opted out via "Don't show again" (see + * {@link optOutOfInstallPrompt}). A plain dismissal doesn't opt out, so the + * offer can resurface on a later activation. */ - async maybePromptInstall(): Promise { - if (this.stateStorage.get("databricks.aitools.hideInstallPrompt")) { - return; - } + get shouldPromptInstall(): boolean { + return !this.stateStorage.get("databricks.aitools.hideInstallPrompt"); + } - const install = "Install AI tools"; - const dontShowAgain = "Don't show again"; - const choice = await window.showInformationMessage( - "Install Databricks AI tools?", - { - modal: true, - detail: "Get skills and plugins so your coding agents work effectively with Databricks. You can also install them later from the Databricks configuration panel.", - }, - install, - dontShowAgain - ); - if (choice === dontShowAgain) { - // The user declined and asked not to be prompted again. - await this.stateStorage.set( - "databricks.aitools.hideInstallPrompt", - true - ); - return; - } - 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 - // "initModal" source so telemetry can distinguish first-load prompt - // installs from manual side-pane installs. - await commands.executeCommand( - "databricks.aitools.install", - "initModal" + /** + * Record that the user opted out of the install prompt ("Don't show again"), + * so it won't be offered again. Called by {@link AiToolsCommands} when the + * prompt is declined with the opt-out affordance. + */ + async optOutOfInstallPrompt(): Promise { + await this.stateStorage.set( + "databricks.aitools.hideInstallPrompt", + true ); } @@ -433,8 +414,14 @@ export class AiToolsManager implements Disposable { } /** - * Install AI tools for the given scope, showing progress. Re-detects the - * install state and refreshes the update status afterwards. + * Install AI tools for the given scope. Re-detects the install state and + * refreshes the update status afterwards, and rethrows any {@link + * ProcessError} for {@link AiToolsCommands} to surface — the reconciliation + * still runs first (via `finally`), since a failed install often still + * landed some tools. + * + * `token` is threaded from the caller's progress notification so the CLI run + * is cancellable. * * In Cursor, selecting the Cursor agent means "install the Databricks * marketplace plugin" (a superset of the Cursor skills), not "install the @@ -446,7 +433,8 @@ export class AiToolsManager implements Disposable { async install( scope: AiToolsScope, source?: AiToolsInstallSource, - agents?: string[] + agents?: string[], + token?: CancellationToken ): Promise { let cliAgents = agents; let cursorPlugin = false; @@ -479,19 +467,11 @@ export class AiToolsManager implements Disposable { 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, - cliAgents - ) + await this.cli.aitoolsInstall( + scope, + this.cwdForScope(scope), + token, + cliAgents ); recordEvent({ success: true, @@ -508,49 +488,38 @@ export class AiToolsManager implements Disposable { agents: cliAgents, cursorPlugin, }); - if (e instanceof ProcessError) { - e.showErrorMessage( - "Failed to install Databricks AI tools.", - "databricks.logs.show" - ); - } else { - throw e; - } - // Fall through: a failed install often still installed some tools - // (e.g. one agent's CLI was missing), so refresh the panel to + throw e; + } finally { + // Always reconcile: a failed install often still installed some + // tools (e.g. one agent's CLI was missing), so refresh the panel to // reflect whatever actually landed rather than leaving it stale. + await this.detectInstall(); + await this.resolveInstalled(); } - - await this.detectInstall(); - await this.resolveInstalled(); } /** - * Install a single coding agent into the current install scope, showing - * progress. Used by the per-agent "install" button in the Agents list to - * add an agent that wasn't installed alongside the others. Re-resolves the - * agent statuses afterwards so the row flips to "installed". + * Install a single coding agent into the current install scope. Used by the + * per-agent "install" button in the Agents list to add an agent that wasn't + * installed alongside the others. Re-resolves the agent statuses afterwards + * (even on failure) so the row reflects the real CLI state, and rethrows any + * error for {@link AiToolsCommands} to surface. */ - async installAgent(agentId: string): Promise { + async installAgent( + agentId: string, + token?: CancellationToken + ): Promise { const scope = this._installLocation; if (scope === undefined) { return; } const recordEvent = this.telemetry.start(Events.AITOOLS_INSTALL); try { - await window.withProgress( - { - location: ProgressLocation.Notification, - title: "Installing Databricks AI tools agent", - cancellable: true, - }, - (_progress, token) => - this.cli.aitoolsInstall( - scope, - this.cwdForScope(scope), - token, - [agentId] - ) + await this.cli.aitoolsInstall( + scope, + this.cwdForScope(scope), + token, + [agentId] ); recordEvent({ success: true, @@ -565,65 +534,47 @@ export class AiToolsManager implements Disposable { source: "sidePane", agents: [agentId], }); - if (e instanceof ProcessError) { - e.showErrorMessage( - "Failed to install Databricks AI tools agent.", - "databricks.logs.show" - ); - } else { - throw e; - } - // Fall through to refresh the panel: the install may have partially - // succeeded, so reconcile the row with the real CLI state. + throw e; + } finally { + // Reconcile the row even on failure: the install may have partially + // succeeded. + await this.resolveInstalled(); } - await this.resolveInstalled(); } /** - * Uninstall AI tools for the current install scope, showing progress. - * Re-detects the install state afterwards. + * Uninstall AI tools for the current install scope. Re-detects the install + * state afterwards (even on failure, since a failed uninstall may have + * removed some tools) and rethrows any error for {@link AiToolsCommands} to + * surface. */ - async uninstall(): Promise { + async uninstall(token?: CancellationToken): 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 - ) + await 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.", - "databricks.logs.show" - ); - } else { - throw e; - } - // Fall through: a failed uninstall may have removed some tools, so - // re-detect to reflect the real state rather than leaving it stale. + throw e; + } finally { + await this.detectInstall(); } - await this.detectInstall(); } /** - * Update AI tools for the current install scope, showing progress. + * Update AI tools for the current install scope. Reconciles the cached + * update status with the real CLI state afterwards (even on failure) and + * rethrows any error for {@link AiToolsCommands} to surface. */ - async update(): Promise { + async update(token?: CancellationToken): Promise { const scope = this._installLocation; if (scope === undefined) { return; @@ -632,30 +583,11 @@ export class AiToolsManager implements Disposable { 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 - ) - ); + await 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.", - "databricks.logs.show" - ); - } else { - throw e; - } + throw e; } finally { // Always reconcile the cached update status with the actual CLI // state, even if the update reported an error (it may have diff --git a/packages/databricks-vscode/src/cli/CliWrapper.test.ts b/packages/databricks-vscode/src/cli/CliWrapper.test.ts index eec5d3188..c84c52c7e 100644 --- a/packages/databricks-vscode/src/cli/CliWrapper.test.ts +++ b/packages/databricks-vscode/src/cli/CliWrapper.test.ts @@ -1,12 +1,17 @@ import * as assert from "assert"; -import {Uri} from "vscode"; +import {commands, Uri, window} from "vscode"; 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, mkdtemp, rm} from "node:fs/promises"; import {when, spy, reset, instance, mock} from "ts-mockito"; -import {cancellableExecFile, CliWrapper, waitForProcess} from "./CliWrapper"; +import { + cancellableExecFile, + CliWrapper, + ProcessError, + waitForProcess, +} from "./CliWrapper"; import path from "node:path"; import os from "node:os"; import crypto from "node:crypto"; @@ -73,6 +78,18 @@ describe(__filename, function () { assert.ok(typeof skill.name === "string"); assert.ok(typeof skill.latest_version === "string"); assert.ok(typeof skill.installed === "object"); + + // It also reports the coding agents it knows about, each with a + // display name and detection/management flags; the agent picker + // relies on these fields. + assert.ok(Array.isArray(result.agents)); + assert.ok(result.agents.length > 0); + const agent = result.agents[0]; + assert.ok(typeof agent.name === "string"); + assert.ok(typeof agent.display_name === "string"); + assert.ok(typeof agent.managed === "boolean"); + assert.ok(typeof agent.detected === "boolean"); + assert.ok(typeof agent.installed === "object"); } finally { await rm(tmpDir, {recursive: true, force: true}); } @@ -377,3 +394,62 @@ describe("waitForProcess", () => { assert.equal(stderr, `{"error": "nooo"}`); }); }); + +describe("ProcessError.showErrorMessage", () => { + let originalShowError: typeof window.showErrorMessage; + let originalExecuteCommand: typeof commands.executeCommand; + let executed: string[]; + + beforeEach(() => { + executed = []; + originalShowError = window.showErrorMessage; + // Resolve as if the user clicked the primary action button (the last + // vararg), so both the "Show Logs" and "Assign Values" branches fire. + (window as any).showErrorMessage = async ( + _message: string, + ...items: string[] + ) => items[items.length - 1]; + originalExecuteCommand = commands.executeCommand; + (commands as any).executeCommand = (command: string) => { + executed.push(command); + }; + }); + + afterEach(() => { + (window as any).showErrorMessage = originalShowError; + (commands as any).executeCommand = originalExecuteCommand; + }); + + // `showErrorMessage` handles the toast promise with `.then` (fire and + // forget), so a microtask tick is needed before the executeCommand runs. + async function flush() { + await new Promise((resolve) => setTimeout(resolve, 0)); + } + + it("opens the bundle logs channel by default", async () => { + new ProcessError("boom", 1).showErrorMessage("Prefix."); + await flush(); + assert.deepStrictEqual(executed, ["databricks.bundle.showLogs"]); + }); + + it("opens the given logs channel when one is passed", async () => { + new ProcessError("boom", 1).showErrorMessage( + "Prefix.", + "databricks.logs.show" + ); + await flush(); + assert.deepStrictEqual(executed, ["databricks.logs.show"]); + }); + + it("ignores the logsCommand for the missing-variable path", async () => { + // The "no value assigned to required variable" branch has its own + // fixed set of commands and never consults logsCommand. + new ProcessError( + "no value assigned to required variable foo", + 1 + ).showErrorMessage("Prefix.", "databricks.logs.show"); + await flush(); + assert.ok(!executed.includes("databricks.logs.show")); + assert.ok(executed.includes("databricks.bundle.showLogs")); + }); +}); diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index 23c8f0b23..7d1ac3eaa 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -217,7 +217,7 @@ export async function activate( // Detect install state on activation and, if installed, auto-apply any // available update; otherwise prompt the user (once) to install the tools. // Non-blocking so it doesn't delay activation. - aiToolsManager.initialize(); + aiToolsCommands.initializeCommand()(); // Developer-only "Reset state" command (multi-select of persisted state // keys). Gated to development builds so it isn't exposed to end users; the diff --git a/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.test.ts b/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.test.ts index 3a4a643ee..061383f86 100644 --- a/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.test.ts +++ b/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.test.ts @@ -10,6 +10,7 @@ import { } from "../../aitools/AiToolsManager"; import {resolveProviderResult} from "../../test/utils"; import {AiToolsComponent} from "./AiToolsComponent"; +import {HostUtils} from "../../utils"; function createManager( installLocation: AiToolsInstallLocation, @@ -359,5 +360,61 @@ describe(__filename, () => { {id: "AITOOLS.agent.codex"}, ]); }); + + describe("in Cursor", () => { + let originalIsCursor: typeof HostUtils.isCursor; + + beforeEach(() => { + originalIsCursor = HostUtils.isCursor; + (HostUtils as any).isCursor = () => true; + }); + + afterEach(() => { + (HostUtils as any).isCursor = originalIsCursor; + }); + + it("hides the Cursor agent row (it is managed via the marketplace plugin)", async () => { + const manager = createManager( + "project", + "upToDate", + "0.2.9", + undefined, + [ + agent("claude", "Claude Code", "1.2.0"), + agent("cursor", "Cursor", "0.3.0"), + ] + ); + const rows = await getChildrenOf(manager, { + label: "Agents", + id: "AITOOLS.agents", + }); + assert.deepStrictEqual( + rows.map((r) => r.id), + ["AITOOLS.agent.claude"] + ); + }); + + it("excludes the hidden Cursor agent from the installed count", async () => { + const manager = createManager( + "project", + "upToDate", + "0.2.9", + undefined, + [ + agent("claude", "Claude Code", "1.2.0"), + // Installed, but hidden in Cursor -> must not be counted. + agent("cursor", "Cursor", "0.3.0"), + ] + ); + const children = await getChildrenOf(manager, { + label: "AI tools", + id: "AITOOLS", + }); + const agentsNode = children.find( + (c) => c.id === "AITOOLS.agents" + ); + assert.strictEqual(agentsNode?.description, "1 installed"); + }); + }); }); }); diff --git a/packages/databricks-vscode/src/utils/hostUtils.test.ts b/packages/databricks-vscode/src/utils/hostUtils.test.ts new file mode 100644 index 000000000..919d8e6ea --- /dev/null +++ b/packages/databricks-vscode/src/utils/hostUtils.test.ts @@ -0,0 +1,39 @@ +import {env} from "vscode"; +import assert from "assert"; +import {isCursor} from "./hostUtils"; + +describe(__filename, () => { + let originalAppName: PropertyDescriptor | undefined; + + function stubAppName(value: string) { + Object.defineProperty(env, "appName", { + value, + configurable: true, + }); + } + + beforeEach(() => { + originalAppName = Object.getOwnPropertyDescriptor(env, "appName"); + }); + + afterEach(() => { + if (originalAppName) { + Object.defineProperty(env, "appName", originalAppName); + } + }); + + it("is true when the app name is Cursor", () => { + stubAppName("Cursor"); + assert.strictEqual(isCursor(), true); + }); + + it("matches case-insensitively and as a substring", () => { + stubAppName("cursor nightly"); + assert.strictEqual(isCursor(), true); + }); + + it("is false for plain VS Code", () => { + stubAppName("Visual Studio Code"); + assert.strictEqual(isCursor(), false); + }); +}); diff --git a/packages/databricks-vscode/src/vscode-objs/StateResetCommand.test.ts b/packages/databricks-vscode/src/vscode-objs/StateResetCommand.test.ts new file mode 100644 index 000000000..cb55078a7 --- /dev/null +++ b/packages/databricks-vscode/src/vscode-objs/StateResetCommand.test.ts @@ -0,0 +1,130 @@ +import assert from "assert"; +import {QuickPickItem, commands, window} from "vscode"; +import {anything, capture, instance, mock, verify, when} from "ts-mockito"; +import {StateStorage, StorageKey} from "./StateStorage"; +import {StateResetCommand} from "./StateResetCommand"; + +describe(__filename, () => { + let mockStorage: StateStorage; + let originalShowQuickPick: typeof window.showQuickPick; + let originalShowInfo: typeof window.showInformationMessage; + let originalExecuteCommand: typeof commands.executeCommand; + let executed: string[]; + + // The picker items offered to the user, captured from the last + // showQuickPick call so tests can inspect ordering/labels. + let offeredItems: readonly QuickPickItem[]; + // Behavior applied to the (single) showQuickPick call. + let pickBehavior: ( + items: readonly QuickPickItem[] + ) => QuickPickItem[] | undefined; + // Behavior applied to the reload prompt. + let reloadChoice: string | undefined; + + beforeEach(() => { + mockStorage = mock(StateStorage); + when(mockStorage.storageKeys).thenReturn([ + { + key: "databricks.bundle.target" as StorageKey, + location: "workspace", + }, + { + key: "databricks.aitools.hideInstallPrompt" as StorageKey, + location: "global", + }, + ]); + when(mockStorage.reset(anything())).thenResolve(); + + offeredItems = []; + pickBehavior = () => undefined; + reloadChoice = undefined; + executed = []; + + originalShowQuickPick = window.showQuickPick; + (window as any).showQuickPick = async (items: QuickPickItem[]) => { + offeredItems = items; + return pickBehavior(items); + }; + originalShowInfo = window.showInformationMessage; + (window as any).showInformationMessage = async () => reloadChoice; + originalExecuteCommand = commands.executeCommand; + (commands as any).executeCommand = async (command: string) => { + executed.push(command); + }; + }); + + afterEach(() => { + (window as any).showQuickPick = originalShowQuickPick; + (window as any).showInformationMessage = originalShowInfo; + (commands as any).executeCommand = originalExecuteCommand; + }); + + function createCommand() { + return new StateResetCommand(instance(mockStorage)); + } + + it("offers every storage key, sorted by label, with its location", async () => { + await createCommand().resetCommand()(); + + assert.deepStrictEqual( + offeredItems.map((i) => i.label), + ["databricks.aitools.hideInstallPrompt", "databricks.bundle.target"] + ); + assert.deepStrictEqual( + offeredItems.map((i) => i.description), + ["global", "workspace"] + ); + }); + + it("resets each selected key and offers a window reload", async () => { + pickBehavior = (items) => + items.filter( + (i) => i.label === "databricks.aitools.hideInstallPrompt" + ); + reloadChoice = "Reload Window"; + + await createCommand().resetCommand()(); + + // `reset` is generic; cast for ts-mockito's capture overload. + const [resetKey] = capture(mockStorage.reset as any).last(); + assert.strictEqual(resetKey, "databricks.aitools.hideInstallPrompt"); + verify(mockStorage.reset(anything())).once(); + assert.deepStrictEqual(executed, ["workbench.action.reloadWindow"]); + }); + + it("resets all selected keys", async () => { + pickBehavior = (items) => [...items]; + + await createCommand().resetCommand()(); + + verify(mockStorage.reset(anything())).twice(); + }); + + it("does not reload when the reload prompt is dismissed", async () => { + pickBehavior = (items) => [items[0]]; + reloadChoice = undefined; + + await createCommand().resetCommand()(); + + verify(mockStorage.reset(anything())).once(); + assert.deepStrictEqual(executed, []); + }); + + it("does nothing when the picker is dismissed", async () => { + pickBehavior = () => undefined; + + await createCommand().resetCommand()(); + + verify(mockStorage.reset(anything())).never(); + assert.deepStrictEqual(executed, []); + }); + + it("does nothing when the selection is empty", async () => { + pickBehavior = () => []; + + await createCommand().resetCommand()(); + + verify(mockStorage.reset(anything())).never(); + assert.deepStrictEqual(executed, []); + }); +}); diff --git a/packages/databricks-vscode/src/vscode-objs/WorkspaceFolderManager.test.ts b/packages/databricks-vscode/src/vscode-objs/WorkspaceFolderManager.test.ts index a49f89738..009e6a9d1 100644 --- a/packages/databricks-vscode/src/vscode-objs/WorkspaceFolderManager.test.ts +++ b/packages/databricks-vscode/src/vscode-objs/WorkspaceFolderManager.test.ts @@ -2,7 +2,7 @@ import {workspace} from "vscode"; import {CustomWhenContext} from "./CustomWhenContext"; import {StateStorage} from "./StateStorage"; import {WorkspaceFolderManager} from "./WorkspaceFolderManager"; -import {instance, mock, when} from "ts-mockito"; +import {anything, capture, instance, mock, verify, when} from "ts-mockito"; import assert from "node:assert"; import path from "node:path"; @@ -67,4 +67,60 @@ describe(__filename, () => { workspaceFolder.uri ); }); + + describe("with no folder open", () => { + let originalFolders: PropertyDescriptor | undefined; + + beforeEach(() => { + // Simulate a folderless window: `workspace.workspaceFolders` is + // undefined, so the manager has no active project/workspace folder. + originalFolders = Object.getOwnPropertyDescriptor( + workspace, + "workspaceFolders" + ); + Object.defineProperty(workspace, "workspaceFolders", { + value: undefined, + configurable: true, + }); + }); + + afterEach(() => { + if (originalFolders) { + Object.defineProperty( + workspace, + "workspaceFolders", + originalFolders + ); + } + }); + + it("does not throw during construction and reports the file as outside the active workspace", () => { + // Regression guard: `setIsActiveFileInActiveProject` runs from the + // constructor and must use the non-throwing private field rather + // than the `activeProjectUri` getter, which throws when no folder is + // open. + const stateStorage = mock(); + const customWhenContext = mock(CustomWhenContext); + + const manager = new WorkspaceFolderManager( + instance(customWhenContext), + instance(stateStorage) + ); + + // The when-context is set to false (no folder -> nothing can be in + // the active workspace), and the getter still throws as designed. + verify( + customWhenContext.setIsActiveFileInActiveWorkspace(false) + ).once(); + const [value] = capture( + customWhenContext.setIsActiveFileInActiveWorkspace + ).last(); + assert.strictEqual(value, false); + assert.throws(() => manager.activeProjectUri); + assert.throws(() => manager.activeWorkspaceFolder); + verify( + customWhenContext.setIsActiveFileInActiveWorkspace(anything()) + ).once(); + }); + }); }); From cf3b8c9954f4ed5cd6ca1c829e72161e7c9c5873 Mon Sep 17 00:00:00 2001 From: Russell Clarey Date: Tue, 4 Aug 2026 13:16:58 +0200 Subject: [PATCH 5/6] refactor tests --- .../src/aitools/AiToolsCommands.test.ts | 255 +++++----- .../src/aitools/AiToolsCommands.ts | 43 +- .../src/aitools/AiToolsManager.test.ts | 451 +++++++++--------- .../src/aitools/AiToolsManager.ts | 157 +++--- .../src/aitools/AiToolsModel.test.ts | 78 +++ .../src/aitools/AiToolsModel.ts | 91 ++++ .../src/cli/CliWrapper.test.ts | 10 +- .../databricks-vscode/src/cli/CliWrapper.ts | 8 +- packages/databricks-vscode/src/extension.ts | 7 +- .../AiToolsComponent.test.ts | 84 ++-- .../ui/configuration-view/AiToolsComponent.ts | 15 +- .../ConfigurationDataProvider.ts | 2 +- .../src/utils/hostUtils.test.ts | 21 +- .../databricks-vscode/src/utils/hostUtils.ts | 7 +- .../src/vscode-objs/CustomWhenContext.ts | 16 + .../src/vscode-objs/StateStorage.ts | 4 +- 16 files changed, 715 insertions(+), 534 deletions(-) create mode 100644 packages/databricks-vscode/src/aitools/AiToolsModel.test.ts create mode 100644 packages/databricks-vscode/src/aitools/AiToolsModel.ts diff --git a/packages/databricks-vscode/src/aitools/AiToolsCommands.test.ts b/packages/databricks-vscode/src/aitools/AiToolsCommands.test.ts index 263c786d5..4d2621b0d 100644 --- a/packages/databricks-vscode/src/aitools/AiToolsCommands.test.ts +++ b/packages/databricks-vscode/src/aitools/AiToolsCommands.test.ts @@ -3,19 +3,17 @@ import assert from "assert"; import { CancellationToken, + MessageOptions, Progress, + ProgressOptions, QuickPick, QuickPickItem, - window, } from "vscode"; import {anything, capture, instance, mock, verify, when} from "ts-mockito"; import {ProcessError} from "../cli/CliWrapper"; -import { - AiToolsAgentStatus, - AiToolsManager, - CURSOR_AGENT_ID, -} from "./AiToolsManager"; -import {AiToolsCommands} from "./AiToolsCommands"; +import {AiToolsManager, CURSOR_AGENT_ID} from "./AiToolsManager"; +import {AiToolsAgentStatus} from "./AiToolsModel"; +import {AiToolsCommands, AiToolsPrompter} from "./AiToolsCommands"; import {HostUtils} from "../utils"; /** @@ -80,82 +78,119 @@ function agent( }; } -describe(__filename, () => { - let mockManager: AiToolsManager; - let originalCreateQuickPick: typeof window.createQuickPick; - let originalWithProgress: typeof window.withProgress; - let originalIsCursor: typeof HostUtils.isCursor; - let quickPicks: FakeQuickPick[]; - // Behavior applied to each createQuickPick call, in order. - let behaviors: Array< +// A cancellation token whose `isCancellationRequested` is false; enough to +// stand in for the token withProgress would hand the task. +const fakeToken = { + isCancellationRequested: false, + onCancellationRequested: () => ({dispose() {}}), +} as unknown as CancellationToken; + +/** + * A scriptable {@link AiToolsPrompter} replacing the VS Code `window` seam. + * `quickPickBehaviors` drives each createQuickPick call in order (the FakeQuickPick + * accept/dismiss script); `messageResponses` supplies the string each + * info/warning message resolves to. Records every created pick and the messages + * shown so assertions can inspect them. + */ +class FakePrompter implements AiToolsPrompter { + readonly quickPicks: FakeQuickPick[] = []; + readonly quickPickBehaviors: Array< ( pick: FakeQuickPick ) => {selected: readonly QuickPickItem[]} | "dismiss" - >; - // A cancellation token whose `isCancellationRequested` is false; enough to - // stand in for the token withProgress would hand the task. - const fakeToken = { - isCancellationRequested: false, - onCancellationRequested: () => ({dispose() {}}), - } as unknown as CancellationToken; - - function stubIsCursor(value: boolean) { - (HostUtils as any).isCursor = () => value; + > = []; + // Response for the next info/warning message; shift()ed per call. + readonly messageResponses: Array = []; + readonly shownMessages: Array<{message: string; items: string[]}> = []; + + // Run the progress task synchronously with a non-cancelled token, skipping + // the real notification UI. + withProgress( + _options: ProgressOptions, + task: ( + progress: Progress<{message?: string; increment?: number}>, + token: CancellationToken + ) => Thenable + ): Thenable { + return Promise.resolve(task({report() {}}, fakeToken)); + } + + showInformationMessage( + message: string, + _options: MessageOptions, + ...items: string[] + ): Thenable { + this.shownMessages.push({message, items}); + return Promise.resolve(this.messageResponses.shift()); + } + + showWarningMessage( + message: string, + _options: MessageOptions, + ...items: string[] + ): Thenable { + this.shownMessages.push({message, items}); + return Promise.resolve(this.messageResponses.shift()); + } + + createQuickPick(): QuickPick { + const behavior = + this.quickPickBehaviors.shift() ?? (() => "dismiss" as const); + const pick = new FakeQuickPick(behavior); + this.quickPicks.push(pick); + return pick as unknown as QuickPick; } +} + +/** Select the scope item matching `scope` in the scope picker. */ +function selectScope(scope: "project" | "global") { + return (pick: FakeQuickPick) => ({ + selected: [ + pick.items.find((i) => (i as any).scope === scope) as QuickPickItem, + ], + }); +} + +function stubIsCursor(value: boolean) { + (HostUtils as any).isCursor = () => value; +} + +// Build an AiToolsCommands wired to a mock manager and a scriptable prompter. +// The manager defaults to a folder being open and no agents reported; tests +// override via when(...) on the returned mockManager. quickPicks/behaviors are +// the prompter's arrays, surfaced for convenience. +function setup() { + const mockManager = mock(AiToolsManager); + when(mockManager.hasProjectFolder).thenReturn(true); + when(mockManager.listAgents(anything())).thenResolve([]); + + const prompter = new FakePrompter(); + + return { + mockManager, + prompter, + quickPicks: prompter.quickPicks, + behaviors: prompter.quickPickBehaviors, + commands: new AiToolsCommands(instance(mockManager), prompter), + }; +} + +describe(__filename, () => { + let originalIsCursor: typeof HostUtils.isCursor; beforeEach(() => { - mockManager = mock(AiToolsManager); - when(mockManager.hasProjectFolder).thenReturn(true); - when(mockManager.listAgents(anything())).thenResolve([]); - - quickPicks = []; - behaviors = []; - originalCreateQuickPick = window.createQuickPick; - (window as any).createQuickPick = () => { - const behavior = behaviors.shift() ?? (() => "dismiss" as const); - const pick = new FakeQuickPick(behavior); - quickPicks.push(pick); - return pick as unknown as QuickPick; - }; - - // Run the progress task synchronously with a non-cancelled token, - // skipping the real notification UI. - originalWithProgress = window.withProgress; - (window as any).withProgress = ( - _options: unknown, - task: ( - progress: Progress, - token: CancellationToken - ) => Thenable - ) => task({report() {}}, fakeToken); - - // Default to plain VS Code; Cursor-specific tests opt in. + // Default to plain VS Code; Cursor-specific tests opt in via + // stubIsCursor(true). originalIsCursor = HostUtils.isCursor; stubIsCursor(false); }); afterEach(() => { - (window as any).createQuickPick = originalCreateQuickPick; - (window as any).withProgress = originalWithProgress; (HostUtils as any).isCursor = originalIsCursor; }); - function createCommands() { - return new AiToolsCommands(instance(mockManager)); - } - - /** Select the scope item matching `scope` in the scope picker. */ - function selectScope(scope: "project" | "global") { - return (pick: FakeQuickPick) => ({ - selected: [ - pick.items.find( - (i) => (i as any).scope === scope - ) as QuickPickItem, - ], - }); - } - it("shows the agent picker after the scope picker and passes the selection to install", async () => { + const {commands, mockManager, behaviors, quickPicks} = setup(); when(mockManager.listAgents("global")).thenResolve([ agent("claude-code", "Claude Code", true), agent("cursor", "Cursor", false), @@ -167,7 +202,7 @@ describe(__filename, () => { selected: pick.selectedItems, })); - await createCommands().installCommand()("sidePane"); + await commands.installCommand()("sidePane"); // The agent picker is the second QuickPick created (after scope). assert.strictEqual(quickPicks.length, 2); @@ -186,6 +221,7 @@ describe(__filename, () => { }); it("preselects only the detected agents", async () => { + const {commands, mockManager, behaviors, quickPicks} = setup(); when(mockManager.listAgents("global")).thenResolve([ agent("claude-code", "Claude Code", true), agent("cursor", "Cursor", false), @@ -194,7 +230,7 @@ describe(__filename, () => { behaviors.push(selectScope("global")); behaviors.push((pick) => ({selected: pick.selectedItems})); - await createCommands().installCommand()("sidePane"); + await commands.installCommand()("sidePane"); const agentPick = quickPicks[1]; assert.deepStrictEqual( @@ -212,6 +248,7 @@ describe(__filename, () => { it("starts the Cursor plugin checked and labels it in Cursor", async () => { stubIsCursor(true); + const {commands, mockManager, behaviors, quickPicks} = setup(); when(mockManager.listAgents("global")).thenResolve([ agent("claude-code", "Claude Code", false), agent(CURSOR_AGENT_ID, "Cursor", false), @@ -219,7 +256,7 @@ describe(__filename, () => { behaviors.push(selectScope("global")); behaviors.push((pick) => ({selected: pick.selectedItems})); - await createCommands().installCommand()("sidePane"); + await commands.installCommand()("sidePane"); const agentPick = quickPicks[1]; // Cursor starts checked even though it isn't "detected"; Claude does not. @@ -237,19 +274,21 @@ describe(__filename, () => { it("does not force the Cursor entry checked outside Cursor", async () => { stubIsCursor(false); + const {commands, mockManager, behaviors, quickPicks} = setup(); when(mockManager.listAgents("global")).thenResolve([ agent(CURSOR_AGENT_ID, "Cursor", false), ]); behaviors.push(selectScope("global")); behaviors.push((pick) => ({selected: pick.selectedItems})); - await createCommands().installCommand()("sidePane"); + await commands.installCommand()("sidePane"); const agentPick = quickPicks[1]; assert.deepStrictEqual(agentPick.selectedItems, []); }); it("installs the user's edited selection, not just the detected agents", async () => { + const {commands, mockManager, behaviors} = setup(); when(mockManager.listAgents("global")).thenResolve([ agent("claude-code", "Claude Code", true), agent("cursor", "Cursor", false), @@ -260,17 +299,18 @@ describe(__filename, () => { selected: pick.items.filter((i) => (i as any).agentId === "cursor"), })); - await createCommands().installCommand()("sidePane"); + await commands.installCommand()("sidePane"); const [, , agents] = capture(mockManager.install).last(); assert.deepStrictEqual(agents, ["cursor"]); }); it("skips the agent picker and installs with an empty selection when no agents are reported", async () => { + const {commands, mockManager, behaviors, quickPicks} = setup(); when(mockManager.listAgents("global")).thenResolve([]); behaviors.push(selectScope("global")); - await createCommands().installCommand()("sidePane"); + await commands.installCommand()("sidePane"); // Only the scope picker is created; the agent picker is skipped. assert.strictEqual(quickPicks.length, 1); @@ -280,21 +320,23 @@ describe(__filename, () => { }); it("cancels the install when the agent picker is dismissed", async () => { + const {commands, mockManager, behaviors} = setup(); when(mockManager.listAgents("global")).thenResolve([ agent("claude-code", "Claude Code", true), ]); behaviors.push(selectScope("global")); behaviors.push(() => "dismiss"); - await createCommands().installCommand()("sidePane"); + await commands.installCommand()("sidePane"); verify(mockManager.install(anything(), anything(), anything())).never(); }); it("does not show the agent picker when the scope picker is dismissed", async () => { + const {commands, mockManager, behaviors, quickPicks} = setup(); behaviors.push(() => "dismiss"); - await createCommands().installCommand()("sidePane"); + await commands.installCommand()("sidePane"); assert.strictEqual(quickPicks.length, 1); verify(mockManager.listAgents(anything())).never(); @@ -302,24 +344,26 @@ describe(__filename, () => { }); it("lists agents for the chosen scope", async () => { + const {commands, mockManager, behaviors} = setup(); when(mockManager.listAgents("project")).thenResolve([ agent("claude-code", "Claude Code", true), ]); behaviors.push(selectScope("project")); behaviors.push((pick) => ({selected: pick.selectedItems})); - await createCommands().installCommand()("sidePane"); + await commands.installCommand()("sidePane"); verify(mockManager.listAgents("project")).once(); }); describe("installAgentCommand", () => { it("installs the agent recovered from the tree node id", async () => { + const {commands, mockManager} = setup(); when( mockManager.installAgent(anything(), anything()) ).thenResolve(); - await createCommands().installAgentCommand()({ + await commands.installAgentCommand()({ id: "AITOOLS.agent.codex", }); @@ -328,8 +372,9 @@ describe(__filename, () => { }); it("ignores a node without an agent id", async () => { - await createCommands().installAgentCommand()({id: "AITOOLS"}); - await createCommands().installAgentCommand()(undefined); + const {commands, mockManager} = setup(); + await commands.installAgentCommand()({id: "AITOOLS"}); + await commands.installAgentCommand()(undefined); verify(mockManager.installAgent(anything(), anything())).never(); }); @@ -337,9 +382,10 @@ describe(__filename, () => { describe("addCursorPluginCommand", () => { it("prompts the plugin with the 'pluginButton' source", async () => { + const {commands, mockManager} = setup(); when(mockManager.addCursorPlugin(anything())).thenResolve(); - await createCommands().addCursorPluginCommand()(); + await commands.addCursorPluginCommand()(); const [source] = capture(mockManager.addCursorPlugin).last(); assert.strictEqual(source, "pluginButton"); @@ -348,17 +394,13 @@ describe(__filename, () => { describe("initializeCommand", () => { it("shows the install prompt and installs on accept when not installed", async () => { + const {commands, mockManager, prompter, behaviors} = setup(); when(mockManager.initialize()).thenResolve("promptInstall"); when(mockManager.listAgents("global")).thenResolve([]); - const originalShowInfo = window.showInformationMessage; - (window as any).showInformationMessage = async () => - "Install AI tools"; + prompter.messageResponses.push("Install AI tools"); behaviors.push(selectScope("global")); - try { - await createCommands().initializeCommand()(); - } finally { - (window as any).showInformationMessage = originalShowInfo; - } + + await commands.initializeCommand()(); // Accepting the prompt runs the install flow with the "initModal" // source (never the opt-out). @@ -368,16 +410,12 @@ describe(__filename, () => { }); it("opts out when the prompt is declined with 'Don't show again'", async () => { + const {commands, mockManager, prompter} = setup(); when(mockManager.initialize()).thenResolve("promptInstall"); when(mockManager.optOutOfInstallPrompt()).thenResolve(); - const originalShowInfo = window.showInformationMessage; - (window as any).showInformationMessage = async () => - "Don't show again"; - try { - await createCommands().initializeCommand()(); - } finally { - (window as any).showInformationMessage = originalShowInfo; - } + prompter.messageResponses.push("Don't show again"); + + await commands.initializeCommand()(); verify(mockManager.optOutOfInstallPrompt()).once(); verify( @@ -386,14 +424,11 @@ describe(__filename, () => { }); it("does not opt out or install on a plain dismissal", async () => { + const {commands, mockManager, prompter} = setup(); when(mockManager.initialize()).thenResolve("promptInstall"); - const originalShowInfo = window.showInformationMessage; - (window as any).showInformationMessage = async () => undefined; - try { - await createCommands().initializeCommand()(); - } finally { - (window as any).showInformationMessage = originalShowInfo; - } + prompter.messageResponses.push(undefined); + + await commands.initializeCommand()(); verify(mockManager.optOutOfInstallPrompt()).never(); verify( @@ -402,18 +437,20 @@ describe(__filename, () => { }); it("applies the update when one is available", async () => { + const {commands, mockManager} = setup(); when(mockManager.initialize()).thenResolve("update"); when(mockManager.update(anything())).thenResolve(); - await createCommands().initializeCommand()(); + await commands.initializeCommand()(); verify(mockManager.update(anything())).once(); }); it("does nothing when the action is 'none'", async () => { + const {commands, mockManager} = setup(); when(mockManager.initialize()).thenResolve("none"); - await createCommands().initializeCommand()(); + await commands.initializeCommand()(); verify(mockManager.update(anything())).never(); verify( @@ -424,6 +461,7 @@ describe(__filename, () => { describe("error handling", () => { it("surfaces a ProcessError from update as a toast (does not rethrow)", async () => { + const {commands, mockManager} = setup(); const err = new ProcessError("boom", 1); let shownPrefix: string | undefined; err.showErrorMessage = (prefix?: string) => { @@ -432,7 +470,7 @@ describe(__filename, () => { when(mockManager.update(anything())).thenReject(err); // A ProcessError is caught and rendered, not propagated. - await createCommands().updateCommand()(); + await commands.updateCommand()(); assert.strictEqual( shownPrefix, @@ -441,12 +479,13 @@ describe(__filename, () => { }); it("rethrows a non-ProcessError from update", async () => { + const {commands, mockManager} = setup(); when(mockManager.update(anything())).thenReject( new Error("unexpected") ); await assert.rejects( - () => createCommands().updateCommand()(), + () => commands.updateCommand()(), /unexpected/ ); }); diff --git a/packages/databricks-vscode/src/aitools/AiToolsCommands.ts b/packages/databricks-vscode/src/aitools/AiToolsCommands.ts index 15297aa04..246376ca4 100644 --- a/packages/databricks-vscode/src/aitools/AiToolsCommands.ts +++ b/packages/databricks-vscode/src/aitools/AiToolsCommands.ts @@ -1,15 +1,12 @@ import { - CancellationToken, - Disposable, ProgressLocation, - QuickPickItem, - window, + type CancellationToken, + type Disposable, + type QuickPickItem, + type window, } from "vscode"; -import { - AiToolsAgentStatus, - AiToolsManager, - CURSOR_AGENT_ID, -} from "./AiToolsManager"; +import {AiToolsManager, CURSOR_AGENT_ID} from "./AiToolsManager"; +import {AiToolsAgentStatus} from "./AiToolsModel"; import {AiToolsScope, ProcessError} from "../cli/CliWrapper"; import {AiToolsInstallSource} from "../telemetry/constants"; import {HostUtils} from "../utils"; @@ -22,10 +19,24 @@ interface AgentQuickPickItem extends QuickPickItem { agentId: string; } +/** + * The prompting UI surface {@link AiToolsCommands} needs, behind one seam. The + * real implementation just delegates to `window` + */ +export interface AiToolsPrompter { + withProgress: (typeof window)["withProgress"]; + showInformationMessage: (typeof window)["showInformationMessage"]; + showWarningMessage: (typeof window)["showWarningMessage"]; + createQuickPick: (typeof window)["createQuickPick"]; +} + export class AiToolsCommands implements Disposable { private disposables: Disposable[] = []; - constructor(private readonly aiToolsManager: AiToolsManager) {} + constructor( + private readonly aiToolsManager: AiToolsManager, + private readonly prompter: AiToolsPrompter + ) {} dispose() { this.disposables.forEach((d) => d.dispose()); @@ -44,7 +55,7 @@ export class AiToolsCommands implements Disposable { run: (token: CancellationToken) => Promise ): Promise { try { - await window.withProgress( + await this.prompter.withProgress( { location: ProgressLocation.Notification, title, @@ -88,7 +99,7 @@ export class AiToolsCommands implements Disposable { private async promptInstall(): Promise { const install = "Install AI tools"; const dontShowAgain = "Don't show again"; - const choice = await window.showInformationMessage( + const choice = await this.prompter.showInformationMessage( "Install Databricks AI tools?", { modal: true, @@ -146,7 +157,7 @@ export class AiToolsCommands implements Disposable { */ private pickScope(): Promise { const hasFolder = this.aiToolsManager.hasProjectFolder; - const quickPick = window.createQuickPick(); + const quickPick = this.prompter.createQuickPick(); quickPick.title = "Install Databricks AI tools"; quickPick.placeholder = "Choose where to install the AI tools"; quickPick.items = [ @@ -211,7 +222,7 @@ export class AiToolsCommands implements Disposable { } const inCursor = HostUtils.isCursor(); - const quickPick = window.createQuickPick(); + const quickPick = this.prompter.createQuickPick(); quickPick.title = "Install Databricks AI tools"; quickPick.placeholder = "Choose which coding agents to install for"; quickPick.canSelectMany = true; @@ -284,11 +295,11 @@ export class AiToolsCommands implements Disposable { uninstallCommand() { return async () => { - const location = this.aiToolsManager.state.installLocation; + const location = this.aiToolsManager.model.installLocation; if (location === undefined) { return; } - const confirm = await window.showWarningMessage( + const confirm = await this.prompter.showWarningMessage( `Uninstall Databricks AI tools (${location})?`, {modal: true}, "Uninstall" diff --git a/packages/databricks-vscode/src/aitools/AiToolsManager.test.ts b/packages/databricks-vscode/src/aitools/AiToolsManager.test.ts index 57f81920c..cf153c5d3 100644 --- a/packages/databricks-vscode/src/aitools/AiToolsManager.test.ts +++ b/packages/databricks-vscode/src/aitools/AiToolsManager.test.ts @@ -1,19 +1,9 @@ /* eslint-disable @typescript-eslint/naming-convention */ import assert from "assert"; -import { - anything, - capture, - instance, - mock, - reset, - verify, - when, -} from "ts-mockito"; +import {anything, capture, instance, mock, verify, when} from "ts-mockito"; import {commands, Uri} from "vscode"; -import {mkdtemp, mkdir, writeFile, rm} from "fs/promises"; import path from "path"; -import os from "os"; import { AiToolsAgent, AiToolsListResult, @@ -22,17 +12,42 @@ import { } from "../cli/CliWrapper"; import {StateStorage} from "../vscode-objs/StateStorage"; import {WorkspaceFolderManager} from "../vscode-objs/WorkspaceFolderManager"; +import {CustomWhenContext} from "../vscode-objs/CustomWhenContext"; import {Telemetry} from "../telemetry"; import {Events} from "../telemetry/constants"; -import {AiToolsManager, CURSOR_AGENT_ID} from "./AiToolsManager"; -import {HostUtils} from "../utils"; - -const STATE_FILE_RELATIVE_PATH = path.join( - ".databricks", - "aitools", - "skills", - ".state.json" -); +import { + AiToolsManager, + CURSOR_AGENT_ID, + StateFileLoader, +} from "./AiToolsManager"; +import {FileUtils, HostUtils} from "../utils"; + +// State-file loaders passed to createManager, one per scope, standing in for the +// on-disk `.state.json` read with no real I/O. The manager only cares whether +// the read resolves or rejects (and whether the rejection is ENOENT). + +/** The state file exists and was read successfully -> tools are installed. */ +const loadSuccess: StateFileLoader = async () => undefined; + +/** The state file is absent (ENOENT) -> tools are not installed. */ +const throwNotFound: StateFileLoader = async () => { + const err: NodeJS.ErrnoException = new Error("ENOENT: no such file"); + err.code = "ENOENT"; + throw err; +}; + +/** An unexpected read failure (non-ENOENT), e.g. a permission/IO error. */ +const throwReadError: StateFileLoader = async () => { + const err: NodeJS.ErrnoException = new Error("EACCES: permission denied"); + err.code = "EACCES"; + throw err; +}; + +/** Per-scope state-file loaders; a missing scope defaults to {@link throwNotFound}. */ +interface ScopeLoaders { + project?: StateFileLoader; + global?: StateFileLoader; +} function listResult( skills: Array<{ @@ -52,187 +67,171 @@ function listResult( }; } +// The project root the WorkspaceFolderManager mock reports; the global root is +// the real home dir (the manager derives it from getHomedir()). The injected +// loader maps a state-file path back to its scope by the projectDir prefix. +const projectDir = path.join(path.sep, "tmp", "aitools-proj"); +const homeDir = FileUtils.getHomedir(); + +// Build a manager whose state-file loader resolves per scope. `loaders` is +// read on every load, so tests can reassign `loaders.project` / +// `loaders.global` mid-test to simulate an install/uninstall. A scope with +// no loader defaults to throwNotFound (not installed). +function setup(loaders: ScopeLoaders = {}) { + const loadStateFile: StateFileLoader = (p) => { + const scope = p.startsWith(projectDir) ? "project" : "global"; + return (loaders[scope] ?? throwNotFound)(p); + }; + + const mockCli = mock(CliWrapper); + + const mockWorkspaceFolderManager = mock(WorkspaceFolderManager); + when(mockWorkspaceFolderManager.activeProjectUri).thenReturn( + Uri.file(projectDir) + ); + + const stubStateStorage = { + state: {} as Record, + get(key: string) { + return this.state[key]; + }, + set(key: string, value: any) { + this.state[key] = value; + }, + onDidChange: () => ({dispose() {}}), + }; + const stubTelemetry = { + events: [] as {event: string; props: any}[], + start(event: string) { + return (props: any) => { + this.events.push({event, props}); + }; + }, + recordEvent(event: string, props: any) { + this.events.push({event, props}); + }, + eventsOfType(event: string) { + return this.events.filter((e) => e.event === event); + }, + }; + + return { + mockCli, + mockWorkspaceFolderManager, + stubStateStorage, + stubTelemetry, + manager: new AiToolsManager( + instance(mockCli), + stubStateStorage as unknown as StateStorage, + instance(mockWorkspaceFolderManager), + // A real CustomWhenContext delegates to commands.executeCommand, + // which the when-context test stubs to observe setContext calls. + new CustomWhenContext(), + stubTelemetry as unknown as Telemetry, + loadStateFile + ), + }; +} + +function stubIsCursor(value: boolean) { + (HostUtils as any).isCursor = () => value; +} + 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; let originalIsCursor: typeof HostUtils.isCursor; - // Every telemetry event recorded during a test (via start()'s recorder or a - // direct recordEvent), so assertions can inspect the emitted properties. - let recordedEvents: Array<{event: string; props: any}>; - - function eventsOfType(event: string) { - return recordedEvents.filter((e) => e.event === event); - } - - function stubIsCursor(value: boolean) { - (HostUtils as any).isCursor = () => value; - } - - 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) - ); - // Capture recorded events. start() returns a recorder callback that - // records under the event name; recordEvent records directly. - recordedEvents = []; - telemetry = { - start: (event: string) => (props: any) => { - recordedEvents.push({event, props}); - }, - recordEvent: (event: string, props: any) => { - recordedEvents.push({event, props}); - }, - } as unknown as Telemetry; + beforeEach(() => { // Default to plain VS Code; Cursor-specific tests opt in via // stubIsCursor(true). originalIsCursor = HostUtils.isCursor; stubIsCursor(false); }); - afterEach(async () => { - process.env.HOME = originalHome; + afterEach(() => { (HostUtils as any).isCursor = originalIsCursor; - 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 {manager, stubStateStorage} = setup(); const location = await manager.detectInstall(); assert.strictEqual(location, undefined); assert.strictEqual(manager.isInstalled, false); assert.strictEqual( - storedState["databricks.aitools.installLocation"], + stubStateStorage.get("databricks.aitools.installLocation"), undefined ); }); it("detects a project install", async () => { - await writeStateFile(projectDir); - const manager = createManager(); + const {manager, stubStateStorage} = setup({project: loadSuccess}); const location = await manager.detectInstall(); assert.strictEqual(location, "project"); assert.strictEqual(manager.isInstalled, true); assert.strictEqual( - storedState["databricks.aitools.installLocation"], + stubStateStorage.get("databricks.aitools.installLocation"), "project" ); }); it("detects a global install when only the home state file exists", async () => { - await writeStateFile(homeDir); - const manager = createManager(); + const {manager, stubStateStorage} = setup({global: loadSuccess}); const location = await manager.detectInstall(); assert.strictEqual(location, "global"); assert.strictEqual( - storedState["databricks.aitools.installLocation"], + stubStateStorage.get("databricks.aitools.installLocation"), "global" ); }); it("prefers project over global when both exist", async () => { - await writeStateFile(projectDir); - await writeStateFile(homeDir); - const manager = createManager(); + const {manager} = setup({ + project: loadSuccess, + global: loadSuccess, + }); 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(); + const loaders: ScopeLoaders = {project: loadSuccess}; + const {manager, stubStateStorage} = setup(loaders); assert.strictEqual(await manager.detectInstall(), "project"); - assert.strictEqual(manager.state.detectError ?? false, false); + assert.strictEqual(manager.model.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)); + // Now the state-file read fails unexpectedly (a non-ENOENT error). + loaders.project = throwReadError; 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(manager.model.state.installLocation, "project"); + assert.strictEqual(manager.model.state.detectError, true); assert.strictEqual( - storedState["databricks.aitools.installLocation"], + stubStateStorage.get("databricks.aitools.installLocation"), "project" ); }); it("clears the detect error flag on a subsequent successful detect", async () => { - await writeStateFile(projectDir); - const manager = createManager(); + const loaders: ScopeLoaders = {project: loadSuccess}; + const {manager} = setup(loaders); 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)); + // Trigger an unexpected read error, then recover. + loaders.project = throwReadError; await manager.detectInstall(); - assert.strictEqual(manager.state.detectError, true); + assert.strictEqual(manager.model.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); + // Restore a readable state file; detection should succeed and clear the flag. + loaders.project = loadSuccess; await manager.detectInstall(); - assert.strictEqual(manager.state.detectError, false); - assert.strictEqual(manager.state.installLocation, "project"); + assert.strictEqual(manager.model.state.detectError, false); + assert.strictEqual(manager.model.state.installLocation, "project"); }); it("reports upToDate when all installed skills match latest", async () => { - await writeStateFile(projectDir); + const {manager, mockCli} = setup({project: loadSuccess}); when(mockCli.aitoolsList(anything())).thenResolve( listResult([ { @@ -242,14 +241,13 @@ describe(__filename, () => { }, ]) ); - const manager = createManager(); await manager.detectInstall(); await manager.resolveInstalled(); - assert.strictEqual(manager.state.updateStatus, "upToDate"); + assert.strictEqual(manager.model.state.updateStatus, "upToDate"); }); it("reports updateAvailable when an installed skill is behind latest", async () => { - await writeStateFile(projectDir); + const {manager, mockCli} = setup({project: loadSuccess}); when(mockCli.aitoolsList(anything())).thenResolve( listResult([ { @@ -264,14 +262,13 @@ describe(__filename, () => { }, ]) ); - const manager = createManager(); await manager.detectInstall(); await manager.resolveInstalled(); - assert.strictEqual(manager.state.updateStatus, "updateAvailable"); + assert.strictEqual(manager.model.state.updateStatus, "updateAvailable"); }); it("ignores non-installed skills when computing update status", async () => { - await writeStateFile(projectDir); + const {manager, mockCli} = setup({project: loadSuccess}); when(mockCli.aitoolsList(anything())).thenResolve( listResult([ { @@ -288,39 +285,35 @@ describe(__filename, () => { }, ]) ); - const manager = createManager(); await manager.detectInstall(); await manager.resolveInstalled(); - assert.strictEqual(manager.state.updateStatus, "upToDate"); + assert.strictEqual(manager.model.state.updateStatus, "upToDate"); }); it("reports error when the list command fails", async () => { - await writeStateFile(projectDir); + const {manager, mockCli} = setup({project: loadSuccess}); when(mockCli.aitoolsList(anything())).thenReject(new Error("boom")); - const manager = createManager(); await manager.detectInstall(); await manager.resolveInstalled(); - assert.strictEqual(manager.state.updateStatus, "error"); + assert.strictEqual(manager.model.state.updateStatus, "error"); }); it("returns unknown update status when not installed", async () => { - const manager = createManager(); + const {manager, mockCli} = setup(); await manager.detectInstall(); await manager.resolveInstalled(); - assert.strictEqual(manager.state.updateStatus, "unknown"); + assert.strictEqual(manager.model.state.updateStatus, "unknown"); verify(mockCli.aitoolsList(anything())).never(); }); it("uninstalls for the detected scope and re-detects", async () => { - await writeStateFile(projectDir); + const loaders: ScopeLoaders = {project: loadSuccess}; + const {manager, mockCli, stubStateStorage} = setup(loaders); when( mockCli.aitoolsUninstall("project", anything(), anything()) ).thenCall(async () => { - await rm(path.join(projectDir, STATE_FILE_RELATIVE_PATH), { - force: true, - }); + loaders.project = throwNotFound; }); - const manager = createManager(); await manager.detectInstall(); assert.strictEqual(manager.isInstalled, true); @@ -331,7 +324,7 @@ describe(__filename, () => { ).once(); assert.strictEqual(manager.isInstalled, false); assert.strictEqual( - storedState["databricks.aitools.installLocation"], + stubStateStorage.get("databricks.aitools.installLocation"), undefined ); }); @@ -351,15 +344,13 @@ describe(__filename, () => { } }; try { - await writeStateFile(projectDir); + const loaders: ScopeLoaders = {project: loadSuccess}; + const {manager, mockCli} = setup(loaders); when( mockCli.aitoolsUninstall("project", anything(), anything()) ).thenCall(async () => { - await rm(path.join(projectDir, STATE_FILE_RELATIVE_PATH), { - force: true, - }); + loaders.project = throwNotFound; }); - const manager = createManager(); await manager.detectInstall(); await manager.uninstall(); @@ -374,7 +365,7 @@ describe(__filename, () => { }); it("does not call the CLI when uninstalling with nothing installed", async () => { - const manager = createManager(); + const {manager, mockCli} = setup(); await manager.detectInstall(); await manager.uninstall(); verify( @@ -383,10 +374,12 @@ describe(__filename, () => { }); it("detects install and refreshes status after a global install", async () => { + const loaders: ScopeLoaders = {}; + const {manager, mockCli} = setup(loaders); when( mockCli.aitoolsInstall("global", anything(), anything(), anything()) ).thenCall(async () => { - await writeStateFile(homeDir); + loaders.global = loadSuccess; }); when(mockCli.aitoolsList(anything())).thenResolve( listResult([ @@ -397,15 +390,14 @@ describe(__filename, () => { }, ]) ); - const manager = createManager(); await manager.install("global"); verify( mockCli.aitoolsInstall("global", anything(), anything(), anything()) ).once(); - assert.strictEqual(manager.state.installLocation, "global"); - assert.strictEqual(manager.state.updateStatus, "upToDate"); + assert.strictEqual(manager.model.state.installLocation, "global"); + assert.strictEqual(manager.model.state.updateStatus, "upToDate"); }); describe("Cursor install flow", () => { @@ -435,7 +427,9 @@ describe(__filename, () => { } it("opens the plugin modal and strips cursor from the CLI --agents", async () => { - await writeStateFile(projectDir); + const {manager, mockCli, stubTelemetry} = setup({ + project: loadSuccess, + }); when( mockCli.aitoolsInstall( "project", @@ -453,7 +447,6 @@ describe(__filename, () => { }, ]) ); - const manager = createManager(); await manager.install("project", "sidePane", [ "claude-code", @@ -470,12 +463,14 @@ describe(__filename, () => { // The install event records only the CLI agents (cursor stripped) // and flags the plugin separately. - const [installEvent] = eventsOfType(Events.AITOOLS_INSTALL); + const [installEvent] = stubTelemetry.eventsOfType( + Events.AITOOLS_INSTALL + ); assert.deepStrictEqual(installEvent.props.agents, ["claude-code"]); assert.strictEqual(installEvent.props.cursorPlugin, true); // Prompting the plugin is recorded too, inheriting the install's // source. - const [pluginEvent] = eventsOfType( + const [pluginEvent] = stubTelemetry.eventsOfType( Events.AITOOLS_CURSOR_PLUGIN_PROMPT ); assert.strictEqual(pluginEvent.props.success, true); @@ -483,8 +478,9 @@ describe(__filename, () => { }); it("skips the CLI install when only the Cursor plugin is selected", async () => { - await writeStateFile(projectDir); - const manager = createManager(); + const {manager, mockCli, stubTelemetry} = setup({ + project: loadSuccess, + }); await manager.install("project", "sidePane", [CURSOR_AGENT_ID]); @@ -504,7 +500,9 @@ describe(__filename, () => { ).never(); // The plugin-only install is still recorded, with no CLI agents. - const [installEvent] = eventsOfType(Events.AITOOLS_INSTALL); + const [installEvent] = stubTelemetry.eventsOfType( + Events.AITOOLS_INSTALL + ); assert.ok(installEvent, "expected an install event"); assert.strictEqual(installEvent.props.success, true); assert.deepStrictEqual(installEvent.props.agents, []); @@ -512,7 +510,7 @@ describe(__filename, () => { }); it("does not open the plugin modal when cursor is not selected", async () => { - await writeStateFile(projectDir); + const {manager, mockCli} = setup({project: loadSuccess}); when( mockCli.aitoolsInstall( "project", @@ -530,7 +528,6 @@ describe(__filename, () => { }, ]) ); - const manager = createManager(); await manager.install("project", "sidePane", ["claude-code"]); @@ -553,11 +550,11 @@ describe(__filename, () => { it("records a successful plugin prompt with the given source", async () => { originalExecuteCommand = commands.executeCommand; (commands as any).executeCommand = async () => {}; - const manager = createManager(); + const {manager, stubTelemetry} = setup(); await manager.addCursorPlugin("pluginButton"); - const [pluginEvent] = eventsOfType( + const [pluginEvent] = stubTelemetry.eventsOfType( Events.AITOOLS_CURSOR_PLUGIN_PROMPT ); assert.strictEqual(pluginEvent.props.success, true); @@ -572,21 +569,21 @@ describe(__filename, () => { throw new Error("no marketplace"); } }; - const manager = createManager(); + const {manager, stubTelemetry} = setup(); await manager.addCursorPlugin(); assert.deepStrictEqual( - eventsOfType(Events.AITOOLS_CURSOR_PLUGIN_PROMPT).map( - (e) => e.props.success - ), + stubTelemetry + .eventsOfType(Events.AITOOLS_CURSOR_PLUGIN_PROMPT) + .map((e) => e.props.success), [false] ); }); }); it("installs a single agent into the current scope and refreshes status", async () => { - await writeStateFile(projectDir); + const {manager, mockCli, stubTelemetry} = setup({project: loadSuccess}); when( mockCli.aitoolsInstall( "project", @@ -604,7 +601,6 @@ describe(__filename, () => { }, ]) ); - const manager = createManager(); await manager.detectInstall(); await manager.installAgent("codex"); @@ -615,14 +611,16 @@ describe(__filename, () => { verify(mockCli.aitoolsList(anything())).once(); // The install event records which agent was installed. - const [installEvent] = eventsOfType(Events.AITOOLS_INSTALL); + const [installEvent] = stubTelemetry.eventsOfType( + Events.AITOOLS_INSTALL + ); assert.strictEqual(installEvent.props.success, true); assert.strictEqual(installEvent.props.source, "sidePane"); assert.deepStrictEqual(installEvent.props.agents, ["codex"]); }); it("does not install an agent when nothing is installed", async () => { - const manager = createManager(); + const {manager, mockCli} = setup(); await manager.detectInstall(); await manager.installAgent("codex"); @@ -638,11 +636,11 @@ describe(__filename, () => { }); it("still refreshes the panel when a single-agent install fails", async () => { + const {manager, mockCli} = setup({project: loadSuccess}); // A partial install (e.g. an agent whose CLI is missing) often still // installed some tools, so the panel must reconcile with the real state // rather than staying stale. The error is rethrown for the command layer // to surface. - await writeStateFile(projectDir); when( mockCli.aitoolsInstall( "project", @@ -660,22 +658,23 @@ describe(__filename, () => { }, ]) ); - const manager = createManager(); await manager.detectInstall(); await assert.rejects(() => manager.installAgent("codex"), ProcessError); // resolveInstalled ran despite the failure. verify(mockCli.aitoolsList(anything())).once(); - assert.strictEqual(manager.state.updateStatus, "upToDate"); + assert.strictEqual(manager.model.state.updateStatus, "upToDate"); }); it("still refreshes the panel when the install command fails", async () => { + const loaders: ScopeLoaders = {}; + const {manager, mockCli} = setup(loaders); when( mockCli.aitoolsInstall("global", anything(), anything(), anything()) ).thenCall(async () => { // Simulate a partial install: some tools landed before the failure. - await writeStateFile(homeDir); + loaders.global = loadSuccess; throw new ProcessError("boom", 1); }); when(mockCli.aitoolsList(anything())).thenResolve( @@ -687,7 +686,6 @@ describe(__filename, () => { }, ]) ); - const manager = createManager(); await assert.rejects( () => manager.install("global", "sidePane", ["codex"]), @@ -696,12 +694,12 @@ describe(__filename, () => { // detectInstall + resolveInstalled ran despite the failure, so the row // reflects the tools that actually installed. - assert.strictEqual(manager.state.installLocation, "global"); - assert.strictEqual(manager.state.updateStatus, "upToDate"); + assert.strictEqual(manager.model.state.installLocation, "global"); + assert.strictEqual(manager.model.state.updateStatus, "upToDate"); }); it("refreshes update status to upToDate after a successful update", async () => { - await writeStateFile(projectDir); + const {manager, mockCli} = setup({project: loadSuccess}); when( mockCli.aitoolsUpdate("project", anything(), anything()) ).thenResolve(); @@ -715,17 +713,16 @@ describe(__filename, () => { }, ]) ); - const manager = createManager(); await manager.detectInstall(); await manager.update(); - assert.strictEqual(manager.state.updateStatus, "upToDate"); + assert.strictEqual(manager.model.state.updateStatus, "upToDate"); verify(mockCli.aitoolsList(anything())).once(); }); it("still refreshes update status when the update command fails", async () => { - await writeStateFile(projectDir); + const {manager, mockCli} = setup({project: loadSuccess}); when( mockCli.aitoolsUpdate("project", anything(), anything()) ).thenReject(new ProcessError("boom", 1)); @@ -738,18 +735,17 @@ describe(__filename, () => { }, ]) ); - const manager = createManager(); await manager.detectInstall(); await assert.rejects(() => manager.update(), ProcessError); // The finally block reconciles state even though the update errored. - assert.strictEqual(manager.state.updateStatus, "updateAvailable"); + assert.strictEqual(manager.model.state.updateStatus, "updateAvailable"); verify(mockCli.aitoolsList(anything())).once(); }); it("captures the installed release version from list", async () => { - await writeStateFile(projectDir); + const {manager, mockCli} = setup({project: loadSuccess}); when(mockCli.aitoolsList(anything())).thenResolve({ release: "0.3.1", skills: [ @@ -762,14 +758,14 @@ describe(__filename, () => { ], agents: [], }); - const manager = createManager(); await manager.detectInstall(); await manager.resolveInstalled(); - assert.strictEqual(manager.state.version, "0.3.1"); + assert.strictEqual(manager.model.state.version, "0.3.1"); }); it("clears the version when nothing is installed", async () => { - await writeStateFile(projectDir); + const loaders: ScopeLoaders = {project: loadSuccess}; + const {manager, mockCli} = setup(loaders); when(mockCli.aitoolsList(anything())).thenResolve( listResult([ { @@ -779,32 +775,34 @@ describe(__filename, () => { }, ]) ); - const manager = createManager(); await manager.detectInstall(); await manager.resolveInstalled(); - assert.strictEqual(manager.state.version, "0.2.9"); + assert.strictEqual(manager.model.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, - }); + loaders.project = throwNotFound; await manager.detectInstall(); - assert.strictEqual(manager.state.version, undefined); + assert.strictEqual(manager.model.state.version, undefined); }); describe("no open folder", () => { - beforeEach(() => { - // Mirror WorkspaceFolderManager throwing when no folder is active. - when(mockWorkspaceFolderManager.activeProjectUri).thenThrow( + // Like setup(), but with no active workspace folder: activeProjectUri + // throws, mirroring WorkspaceFolderManager when no folder is active. + function setupNoFolder(loaders: ScopeLoaders = {}) { + const s = setup(loaders); + when(s.mockWorkspaceFolderManager.activeProjectUri).thenThrow( new Error("No active project folder") ); - }); + return s; + } it("reports no project folder", () => { - assert.strictEqual(createManager().hasProjectFolder, false); + assert.strictEqual(setupNoFolder().manager.hasProjectFolder, false); }); it("installs global against the home dir without touching projectRoot", async () => { + const loaders: ScopeLoaders = {}; + const {manager, mockCli} = setupNoFolder(loaders); when( mockCli.aitoolsInstall( "global", @@ -813,7 +811,7 @@ describe(__filename, () => { anything() ) ).thenCall(async () => { - await writeStateFile(homeDir); + loaders.global = loadSuccess; }); when(mockCli.aitoolsList(anything())).thenResolve( listResult([ @@ -824,7 +822,6 @@ describe(__filename, () => { }, ]) ); - const manager = createManager(); // Must not throw even though no folder is open. await manager.install("global"); @@ -837,19 +834,18 @@ describe(__filename, () => { anything() ) ).once(); - assert.strictEqual(manager.state.installLocation, "global"); + assert.strictEqual(manager.model.state.installLocation, "global"); }); it("detects a global install with no folder open", async () => { - await writeStateFile(homeDir); - const manager = createManager(); + const {manager} = setupNoFolder({global: loadSuccess}); assert.strictEqual(await manager.detectInstall(), "global"); }); }); describe("initialize", () => { it("resolves the update status but reports 'update' when installed and behind", async () => { - await writeStateFile(projectDir); + const {manager, mockCli} = setup({project: loadSuccess}); when(mockCli.aitoolsList(anything())).thenResolve( listResult([ { @@ -859,21 +855,23 @@ describe(__filename, () => { }, ]) ); - const manager = createManager(); const action = await manager.initialize(); // initialize resolves status but leaves applying the update to the // caller (AiToolsCommands). assert.strictEqual(action, "update"); - assert.strictEqual(manager.state.updateStatus, "updateAvailable"); + assert.strictEqual( + manager.model.state.updateStatus, + "updateAvailable" + ); verify( mockCli.aitoolsUpdate(anything(), anything(), anything()) ).never(); }); it("reports 'none' when installed and up to date", async () => { - await writeStateFile(projectDir); + const {manager, mockCli} = setup({project: loadSuccess}); when(mockCli.aitoolsList(anything())).thenResolve( listResult([ { @@ -883,32 +881,31 @@ describe(__filename, () => { }, ]) ); - const manager = createManager(); assert.strictEqual(await manager.initialize(), "none"); }); it("reports 'promptInstall' when not installed", async () => { - const manager = createManager(); + const {manager} = setup(); assert.strictEqual(await manager.initialize(), "promptInstall"); }); it("reports 'none' when not installed but opted out", async () => { - storedState["databricks.aitools.hideInstallPrompt"] = true; - const manager = createManager(); + const {manager, stubStateStorage} = setup(); + stubStateStorage.set("databricks.aitools.hideInstallPrompt", true); assert.strictEqual(await manager.initialize(), "none"); }); it("records the opt-out via optOutOfInstallPrompt", async () => { - const manager = createManager(); + const {manager, stubStateStorage} = setup(); assert.strictEqual(manager.shouldPromptInstall, true); await manager.optOutOfInstallPrompt(); assert.strictEqual( - storedState["databricks.aitools.hideInstallPrompt"], + stubStateStorage.get("databricks.aitools.hideInstallPrompt"), true ); assert.strictEqual(manager.shouldPromptInstall, false); diff --git a/packages/databricks-vscode/src/aitools/AiToolsManager.ts b/packages/databricks-vscode/src/aitools/AiToolsManager.ts index 5ba8ee1fc..b3bf8cd66 100644 --- a/packages/databricks-vscode/src/aitools/AiToolsManager.ts +++ b/packages/databricks-vscode/src/aitools/AiToolsManager.ts @@ -1,6 +1,5 @@ -import {readFile} from "fs/promises"; import path from "path"; -import {CancellationToken, commands, Disposable, EventEmitter} from "vscode"; +import {CancellationToken, commands, Disposable} from "vscode"; import {logging} from "@databricks/sdk-experimental"; import { AiToolsAgent, @@ -10,6 +9,7 @@ import { } from "../cli/CliWrapper"; import {StateStorage} from "../vscode-objs/StateStorage"; import {WorkspaceFolderManager} from "../vscode-objs/WorkspaceFolderManager"; +import {CustomWhenContext} from "../vscode-objs/CustomWhenContext"; import {Telemetry} from "../telemetry"; import { AiToolsCursorPluginSource, @@ -18,6 +18,19 @@ import { } from "../telemetry/constants"; import {Loggers} from "../logger"; import {FileUtils, HostUtils} from "../utils"; +import { + AiToolsAgentStatus, + AiToolsInstallLocation, + AiToolsModel, + AiToolsUpdateStatus, +} from "./AiToolsModel"; + +/** + * Reads the aitools state file at `path`, resolving with its contents and + * rejecting with an `ENOENT`-coded error when it is absent (i.e. the same + * contract as `fs/promises` `readFile`) + */ +export type StateFileLoader = (path: string) => Promise; /** Cursor marketplace numeric ID for the Databricks plugin. */ const CURSOR_PLUGIN_ID = "26723531"; @@ -37,18 +50,6 @@ const STATE_FILE_RELATIVE_PATH = path.join( ".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"; - /** * What activation should do after {@link AiToolsManager.initialize} has detected * the install state and (if installed) resolved the update status. The manager @@ -61,29 +62,6 @@ export type AiToolsUpdateStatus = */ export type AiToolsInitAction = "promptInstall" | "update" | "none"; -export interface AiToolsAgentStatus { - displayName: string; - id: string; - type: "plugin" | "skills-only"; - detected: boolean; - version?: string; -} - -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; - agents: AiToolsAgentStatus[]; -} - function computeUpdateStatus( skills: AiToolsSkill[], scope: AiToolsScope @@ -114,40 +92,26 @@ function computeAgentsStatuses( */ 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; - private _agents: AiToolsAgentStatus[] = []; + public readonly model: AiToolsModel; constructor( private readonly cli: CliWrapper, private readonly stateStorage: StateStorage, private readonly workspaceFolderManager: WorkspaceFolderManager, - private readonly telemetry: Telemetry + private readonly customWhenContext: CustomWhenContext, + private readonly telemetry: Telemetry, + private readonly loadStateFile: StateFileLoader ) { - this._installLocation = this.stateStorage.get( - "databricks.aitools.installLocation" + this.model = new AiToolsModel( + 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, - agents: this._agents, - }; - } - get isInstalled(): boolean { - return this._installLocation !== undefined; + return this.model.isInstalled; } /** @@ -188,11 +152,7 @@ export class AiToolsManager implements Disposable { * can (re-)open the plugin modal at any time. */ private refreshCursorPluginContext() { - commands.executeCommand( - "setContext", - "databricks.context.aitools.showCursorPlugin", - HostUtils.isCursor() - ); + this.customWhenContext.setAiToolsShowCursorPlugin(HostUtils.isCursor()); } /** @@ -201,16 +161,12 @@ export class AiToolsManager implements Disposable { * Uninstall appropriately. */ private refreshInstalledContext() { - commands.executeCommand( - "setContext", - "databricks.context.aitools.installed", - this.isInstalled - ); + this.customWhenContext.setAiToolsInstalled(this.isInstalled); } dispose() { this.disposables.forEach((d) => d.dispose()); - this.onDidChangeEmitter.dispose(); + this.model.dispose(); } /** @@ -247,10 +203,10 @@ export class AiToolsManager implements Disposable { private async stateFileExists(scope: AiToolsScope): Promise { try { - await readFile(this.stateFilePath(scope)); + await this.loadStateFile(this.stateFilePath(scope)); return true; - } catch (e: any) { - if (e?.code === "ENOENT") { + } catch (e: unknown) { + if (e instanceof Error && "code" in e && e.code === "ENOENT") { return false; } throw e; @@ -261,7 +217,7 @@ export class AiToolsManager implements Disposable { * 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}. + * location and fires {@link AiToolsModel.onDidChange}. */ async detectInstall(): Promise { let location: AiToolsInstallLocation; @@ -276,8 +232,6 @@ export class AiToolsManager implements Disposable { } 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 @@ -289,23 +243,26 @@ export class AiToolsManager implements Disposable { "Failed to detect Databricks AI tools install state", e ); - this._detectError = true; + this.model.update({detectError: true}); this.refreshInstalledContext(); - this.onDidChangeEmitter.fire(); - return this._installLocation; + return this.model.installLocation; } - this._installLocation = location; await this.stateStorage.set( "databricks.aitools.installLocation", location ); - if (location === undefined) { - this._updateStatus = "unknown"; - this._version = undefined; - } + // Detection succeeded (a definitive present/absent answer), so clear the + // error flag. When nothing is installed, also reset the derived status + // and version. + this.model.update({ + installLocation: location, + detectError: false, + ...(location === undefined + ? {updateStatus: "unknown", version: undefined} + : {}), + }); this.refreshInstalledContext(); - this.onDidChangeEmitter.fire(); return location; } @@ -329,7 +286,9 @@ export class AiToolsManager implements Disposable { return this.shouldPromptInstall ? "promptInstall" : "none"; } await this.resolveInstalled(); - return this._updateStatus === "updateAvailable" ? "update" : "none"; + return this.model.state.updateStatus === "updateAvailable" + ? "update" + : "none"; } catch (e) { logging.NamedLogger.getOrCreate(Loggers.Extension).error( "Failed to initialize Databricks AI tools", @@ -388,29 +347,28 @@ export class AiToolsManager implements Disposable { * source of truth. */ async resolveInstalled(): Promise { - const scope = this._installLocation; + const scope = this.model.installLocation; if (scope === undefined) { - this._updateStatus = "unknown"; - this.onDidChangeEmitter.fire(); + this.model.update({updateStatus: "unknown"}); return; } - this._updateStatus = "checking"; - this.onDidChangeEmitter.fire(); + this.model.update({updateStatus: "checking"}); try { const result = await this.cli.aitoolsList(this.cwdForScope(scope)); - this._version = result.release; - this._updateStatus = computeUpdateStatus(result.skills, scope); - this._agents = computeAgentsStatuses(result.agents, scope); + this.model.update({ + version: result.release, + updateStatus: computeUpdateStatus(result.skills, scope), + agents: computeAgentsStatuses(result.agents, scope), + }); } catch (e) { logging.NamedLogger.getOrCreate(Loggers.Extension).error( "Failed to check for Databricks AI tools updates", e ); - this._updateStatus = "error"; + this.model.update({updateStatus: "error"}); } - this.onDidChangeEmitter.fire(); } /** @@ -509,7 +467,7 @@ export class AiToolsManager implements Disposable { agentId: string, token?: CancellationToken ): Promise { - const scope = this._installLocation; + const scope = this.model.installLocation; if (scope === undefined) { return; } @@ -549,7 +507,7 @@ export class AiToolsManager implements Disposable { * surface. */ async uninstall(token?: CancellationToken): Promise { - const scope = this._installLocation; + const scope = this.model.installLocation; if (scope === undefined) { return; } @@ -575,13 +533,12 @@ export class AiToolsManager implements Disposable { * rethrows any error for {@link AiToolsCommands} to surface. */ async update(token?: CancellationToken): Promise { - const scope = this._installLocation; + const scope = this.model.installLocation; if (scope === undefined) { return; } const recordEvent = this.telemetry.start(Events.AITOOLS_UPDATE); - this._updateStatus = "updating"; - this.onDidChangeEmitter.fire(); + this.model.update({updateStatus: "updating"}); try { await this.cli.aitoolsUpdate(scope, this.cwdForScope(scope), token); recordEvent({success: true, scope}); diff --git a/packages/databricks-vscode/src/aitools/AiToolsModel.test.ts b/packages/databricks-vscode/src/aitools/AiToolsModel.test.ts new file mode 100644 index 000000000..3db3fb4a3 --- /dev/null +++ b/packages/databricks-vscode/src/aitools/AiToolsModel.test.ts @@ -0,0 +1,78 @@ +import assert from "assert"; +import {AiToolsAgentStatus, AiToolsModel} from "./AiToolsModel"; + +function agent(id: string, version?: string): AiToolsAgentStatus { + return { + id, + displayName: id, + type: version !== undefined ? "plugin" : "skills-only", + detected: version !== undefined, + version, + }; +} + +describe(__filename, () => { + it("seeds the state from the given install location", () => { + const model = new AiToolsModel("project"); + assert.deepStrictEqual(model.state, { + installLocation: "project", + updateStatus: "unknown", + version: undefined, + detectError: false, + agents: [], + }); + assert.strictEqual(model.installLocation, "project"); + assert.strictEqual(model.isInstalled, true); + }); + + it("reports not installed when seeded with no location", () => { + const model = new AiToolsModel(undefined); + assert.strictEqual(model.installLocation, undefined); + assert.strictEqual(model.isInstalled, false); + }); + + it("merges a partial patch into the existing state", () => { + const model = new AiToolsModel("project"); + model.update({updateStatus: "checking"}); + model.update({version: "0.2.9", agents: [agent("codex", "0.2.9")]}); + + // The earlier updateStatus survives the later patch (merge, not replace). + assert.strictEqual(model.state.updateStatus, "checking"); + assert.strictEqual(model.state.version, "0.2.9"); + assert.strictEqual(model.state.installLocation, "project"); + assert.deepStrictEqual(model.state.agents, [agent("codex", "0.2.9")]); + }); + + it("fires onDidChange once per update", () => { + const model = new AiToolsModel(undefined); + let fired = 0; + model.onDidChange(() => fired++); + + model.update({installLocation: "global"}); + model.update({updateStatus: "upToDate"}); + + assert.strictEqual(fired, 2); + assert.strictEqual(model.installLocation, "global"); + }); + + it("returns a fresh snapshot object each read", () => { + const model = new AiToolsModel("project"); + assert.notStrictEqual(model.state, model.state); + + // Reassigning a field on the snapshot must not leak back into the model. + const snapshot = model.state; + snapshot.installLocation = "global"; + assert.strictEqual(model.state.installLocation, "project"); + }); + + it("stops firing after dispose", () => { + const model = new AiToolsModel(undefined); + let fired = 0; + model.onDidChange(() => fired++); + model.dispose(); + + // The emitter is disposed; listeners no longer receive events. + model.update({updateStatus: "error"}); + assert.strictEqual(fired, 0); + }); +}); diff --git a/packages/databricks-vscode/src/aitools/AiToolsModel.ts b/packages/databricks-vscode/src/aitools/AiToolsModel.ts new file mode 100644 index 000000000..f85231c0b --- /dev/null +++ b/packages/databricks-vscode/src/aitools/AiToolsModel.ts @@ -0,0 +1,91 @@ +import {Disposable, Event, EventEmitter} from "vscode"; +import {AiToolsScope} from "../cli/CliWrapper"; + +/** 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 AiToolsAgentStatus { + displayName: string; + id: string; + type: "plugin" | "skills-only"; + detected: boolean; + version?: string; +} + +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; + agents: AiToolsAgentStatus[]; +} + +/** + * Holds the observable AI tools state and notifies listeners when it changes. + * + * This is a pull model: the {@link AiToolsManager} owns all the logic and pushes + * state in via {@link update}, while clients (the tree component) listen to + * {@link onDidChange} and read the current snapshot from {@link state}. Keeping + * the state + notification here leaves the manager to focus purely on detecting + * and mutating the install. + */ +export class AiToolsModel implements Disposable { + private _onDidChange: EventEmitter = new EventEmitter(); + readonly onDidChange: Event = this._onDidChange.event; + + private _state: AiToolsState; + + constructor(installLocation: AiToolsInstallLocation) { + this._state = { + installLocation, + updateStatus: "unknown", + version: undefined, + detectError: false, + agents: [], + }; + } + + /** A snapshot of the current state. */ + get state(): AiToolsState { + return {...this._state}; + } + + get isInstalled(): boolean { + return this._state.installLocation !== undefined; + } + + /** The resolved install scope, or undefined when not installed. */ + get installLocation(): AiToolsInstallLocation { + return this._state.installLocation; + } + + /** + * Merge a partial state patch and notify listeners. Every mutation of the + * AI tools state flows through here, so a single {@link onDidChange} fires + * per logical change. + */ + update(patch: Partial): void { + this._state = {...this._state, ...patch}; + this._onDidChange.fire(); + } + + dispose() { + this._onDidChange.dispose(); + } +} diff --git a/packages/databricks-vscode/src/cli/CliWrapper.test.ts b/packages/databricks-vscode/src/cli/CliWrapper.test.ts index e5317d1cb..a4fb670dc 100644 --- a/packages/databricks-vscode/src/cli/CliWrapper.test.ts +++ b/packages/databricks-vscode/src/cli/CliWrapper.test.ts @@ -7,11 +7,11 @@ import {withFile} from "tmp-promise"; import {writeFile, readFile, mkdtemp, rm} from "node:fs/promises"; import {when, spy, reset, instance, mock} from "ts-mockito"; import { - cancellableExecFile, - CliWrapper, - ProcessError, - getSshConnectCommand, - waitForProcess, + cancellableExecFile, + CliWrapper, + ProcessError, + getSshConnectCommand, + waitForProcess, } from "./CliWrapper"; import path from "node:path"; import os from "node:os"; diff --git a/packages/databricks-vscode/src/cli/CliWrapper.ts b/packages/databricks-vscode/src/cli/CliWrapper.ts index 28e96fd86..c2be57d7b 100644 --- a/packages/databricks-vscode/src/cli/CliWrapper.ts +++ b/packages/databricks-vscode/src/cli/CliWrapper.ts @@ -10,7 +10,6 @@ import { Uri, commands, CancellationToken, - env, } from "vscode"; import {workspaceConfigs} from "../vscode-objs/WorkspaceConfigs"; import {promisify} from "node:util"; @@ -18,7 +17,7 @@ import {logging} from "@databricks/sdk-experimental"; import {LoggerManager, Loggers} from "../logger"; import {Context, context} from "@databricks/sdk-experimental/dist/context"; import {Cloud} from "../utils/constants"; -import {EnvVarGenerators, FileUtils, UrlUtils} from "../utils"; +import {EnvVarGenerators, FileUtils, HostUtils, UrlUtils} from "../utils"; import {AuthProvider} from "../configuration/auth/AuthProvider"; import type {AiToolsScope} from "../telemetry/constants"; export type {AiToolsScope}; @@ -117,8 +116,7 @@ export const execFile = async ( * IDE window. Serverless is the default when no cluster is given. * * The --ide flag matches the host editor so the CLI opens the right remote - * window: Cursor identifies itself via env.uriScheme === "cursor", - * everything else (VS Code, Insiders) uses vscode. + * window * * Logging is configured out of band via the DATABRICKS_LOG_* env vars (see * CliWrapper.getSshConnectEnvVars), so we do not pass --log-* flags here. @@ -126,7 +124,7 @@ export const execFile = async ( export function getSshConnectCommand(opts: {compute: SshConnectCompute}): { args: string[]; } { - const ide = env.uriScheme === "cursor" ? "cursor" : "vscode"; + const ide = HostUtils.isCursor() ? "cursor" : "vscode"; const args = ["ssh", "connect", `--ide=${ide}`, "--auto-approve"]; if (opts.compute.type === "cluster") { // Start a stopped single-user cluster when connecting. diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index 04af89262..84f722a45 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -8,6 +8,7 @@ import { window, workspace, } from "vscode"; +import {readFile} from "fs/promises"; import {CliWrapper} from "./cli/CliWrapper"; import {ConnectionCommands} from "./configuration/ConnectionCommands"; import {ConnectionManager} from "./configuration/ConnectionManager"; @@ -328,9 +329,11 @@ export async function activate( cli, stateStorage, workspaceFolderManager, - telemetry + customWhenContext, + telemetry, + (path) => readFile(path) ); - const aiToolsCommands = new AiToolsCommands(aiToolsManager); + const aiToolsCommands = new AiToolsCommands(aiToolsManager, window); context.subscriptions.push( aiToolsManager, aiToolsCommands, diff --git a/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.test.ts b/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.test.ts index 061383f86..300a78173 100644 --- a/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.test.ts +++ b/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.test.ts @@ -5,37 +5,37 @@ import {ThemeIcon} from "vscode"; import { AiToolsAgentStatus, AiToolsInstallLocation, - AiToolsManager, + AiToolsModel, AiToolsUpdateStatus, -} from "../../aitools/AiToolsManager"; +} from "../../aitools/AiToolsModel"; import {resolveProviderResult} from "../../test/utils"; import {AiToolsComponent} from "./AiToolsComponent"; import {HostUtils} from "../../utils"; -function createManager( +function createModel( installLocation: AiToolsInstallLocation, updateStatus: AiToolsUpdateStatus, version?: string, detectError?: boolean, agents: AiToolsAgentStatus[] = [] -): AiToolsManager { +): AiToolsModel { return { state: {installLocation, updateStatus, version, detectError, agents}, onDidChange: () => ({dispose() {}}), - } as unknown as AiToolsManager; + } as unknown as AiToolsModel; } -async function getRoot(manager: AiToolsManager) { - const component = new AiToolsComponent(manager); +async function getRoot(model: AiToolsModel) { + const component = new AiToolsComponent(model); const items = await resolveProviderResult(component.getChildren()); return items ?? []; } async function getChildrenOf( - manager: AiToolsManager, + model: AiToolsModel, parent: {label?: string; id?: string} ) { - const component = new AiToolsComponent(manager); + const component = new AiToolsComponent(model); const items = await resolveProviderResult(component.getChildren(parent)); return items ?? []; } @@ -56,7 +56,7 @@ function agent( describe(__filename, () => { it("renders a setup prompt when not installed", async () => { - const items = await getRoot(createManager(undefined, "unknown")); + const items = await getRoot(createModel(undefined, "unknown")); assert.strictEqual(items.length, 1); const [row] = items; assert.strictEqual( @@ -68,7 +68,7 @@ describe(__filename, () => { it("renders a retry row when detection failed with no cached location", async () => { const items = await getRoot( - createManager(undefined, "unknown", undefined, true) + createModel(undefined, "unknown", undefined, true) ); assert.strictEqual(items.length, 1); const [row] = items; @@ -84,7 +84,7 @@ describe(__filename, () => { 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) + createModel("project", "upToDate", "0.2.9", true) ); const [row] = items; assert.strictEqual(row.label, "AI tools"); @@ -93,7 +93,7 @@ describe(__filename, () => { it("renders the installed version in the subtext for a project install", async () => { const items = await getRoot( - createManager("project", "upToDate", "0.2.9") + createModel("project", "upToDate", "0.2.9") ); assert.strictEqual(items.length, 1); const [row] = items; @@ -107,7 +107,7 @@ describe(__filename, () => { }); it("falls back to 'Up to date' when the version is unknown", async () => { - const items = await getRoot(createManager("project", "upToDate")); + const items = await getRoot(createModel("project", "upToDate")); const [row] = items; assert.ok(String(row.description).includes("Up to date")); // Stable states use the robot icon. @@ -115,7 +115,7 @@ describe(__filename, () => { }); it("renders an update-available row without a click command", async () => { - const items = await getRoot(createManager("global", "updateAvailable")); + const items = await getRoot(createModel("global", "updateAvailable")); const [row] = items; assert.strictEqual( row.contextValue, @@ -129,7 +129,7 @@ describe(__filename, () => { }); it("renders an updating spinner while auto-updating", async () => { - const items = await getRoot(createManager("project", "updating")); + const items = await getRoot(createModel("project", "updating")); const [row] = items; assert.strictEqual( row.contextValue, @@ -140,13 +140,13 @@ describe(__filename, () => { }); it("does not attach a click command to an up-to-date row", async () => { - const items = await getRoot(createManager("project", "upToDate")); + const items = await getRoot(createModel("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 items = await getRoot(createModel("project", "checking")); const [row] = items; assert.strictEqual( row.contextValue, @@ -156,7 +156,7 @@ describe(__filename, () => { }); it("uses the generic installed context value for unknown status", async () => { - const items = await getRoot(createManager("project", "unknown")); + const items = await getRoot(createModel("project", "unknown")); const [row] = items; assert.strictEqual( row.contextValue, @@ -166,7 +166,7 @@ describe(__filename, () => { it("returns nothing for a non-root parent", async () => { const component = new AiToolsComponent( - createManager("project", "upToDate") + createModel("project", "upToDate") ); const children = await resolveProviderResult( component.getChildren({label: "AI tools", id: "unknown"}) @@ -180,7 +180,7 @@ describe(__filename, () => { // register a second element with id "AITOOLS" and throw. it("does not re-emit the root row for a foreign parent when not installed", async () => { const component = new AiToolsComponent( - createManager(undefined, "unknown") + createModel(undefined, "unknown") ); const children = await resolveProviderResult( component.getChildren({label: "Some other node", id: "cluster"}) @@ -190,7 +190,7 @@ describe(__filename, () => { it("does not re-emit the root row for a foreign parent when detection errored", async () => { const component = new AiToolsComponent( - createManager(undefined, "unknown", undefined, true) + createModel(undefined, "unknown", undefined, true) ); const children = await resolveProviderResult( component.getChildren({label: "Some other node", id: "cluster"}) @@ -200,7 +200,7 @@ describe(__filename, () => { describe("agents", () => { it("renders an Agents node summarizing how many are installed", async () => { - const manager = createManager( + const model = createModel( "project", "upToDate", "0.2.9", @@ -211,7 +211,7 @@ describe(__filename, () => { agent("copilot", "GitHub Copilot", "0.5.0"), ] ); - const children = await getChildrenOf(manager, { + const children = await getChildrenOf(model, { label: "AI tools", id: "AITOOLS", }); @@ -223,14 +223,14 @@ describe(__filename, () => { }); it("reports 0 installed when no agents have a version", async () => { - const manager = createManager( + const model = createModel( "project", "upToDate", "0.2.9", undefined, [agent("cursor", "Cursor")] ); - const children = await getChildrenOf(manager, { + const children = await getChildrenOf(model, { label: "AI tools", id: "AITOOLS", }); @@ -239,14 +239,14 @@ describe(__filename, () => { }); it("still renders the Agents node when there are no agents", async () => { - const manager = createManager( + const model = createModel( "project", "upToDate", "0.2.9", undefined, [] ); - const children = await getChildrenOf(manager, { + const children = await getChildrenOf(model, { label: "AI tools", id: "AITOOLS", }); @@ -256,7 +256,7 @@ describe(__filename, () => { }); it("lists each agent with its version under the Agents node", async () => { - const manager = createManager( + const model = createModel( "project", "upToDate", "0.2.9", @@ -266,7 +266,7 @@ describe(__filename, () => { agent("cursor", "Cursor"), ] ); - const rows = await getChildrenOf(manager, { + const rows = await getChildrenOf(model, { label: "Agents", id: "AITOOLS.agents", }); @@ -284,14 +284,14 @@ describe(__filename, () => { }); it("returns no agent rows when the agents list is empty", async () => { - const manager = createManager( + const model = createModel( "project", "upToDate", "0.2.9", undefined, [] ); - const rows = await getChildrenOf(manager, { + const rows = await getChildrenOf(model, { label: "Agents", id: "AITOOLS.agents", }); @@ -299,14 +299,14 @@ describe(__filename, () => { }); it("marks installed agents with a green check and no install affordance", async () => { - const manager = createManager( + const model = createModel( "project", "upToDate", "0.2.9", undefined, [agent("claude", "Claude Code", "1.2.0")] ); - const [row] = await getChildrenOf(manager, { + const [row] = await getChildrenOf(model, { label: "Agents", id: "AITOOLS.agents", }); @@ -320,14 +320,14 @@ describe(__filename, () => { }); it("gives uninstalled agents the install context value used by the inline button", async () => { - const manager = createManager( + const model = createModel( "project", "upToDate", "0.2.9", undefined, [agent("codex", "Codex CLI")] ); - const [row] = await getChildrenOf(manager, { + const [row] = await getChildrenOf(model, { label: "Agents", id: "AITOOLS.agents", }); @@ -339,14 +339,14 @@ describe(__filename, () => { }); it("makes an uninstalled agent row clickable to install it", async () => { - const manager = createManager( + const model = createModel( "project", "upToDate", "0.2.9", undefined, [agent("codex", "Codex CLI")] ); - const [row] = await getChildrenOf(manager, { + const [row] = await getChildrenOf(model, { label: "Agents", id: "AITOOLS.agents", }); @@ -374,7 +374,7 @@ describe(__filename, () => { }); it("hides the Cursor agent row (it is managed via the marketplace plugin)", async () => { - const manager = createManager( + const model = createModel( "project", "upToDate", "0.2.9", @@ -384,7 +384,7 @@ describe(__filename, () => { agent("cursor", "Cursor", "0.3.0"), ] ); - const rows = await getChildrenOf(manager, { + const rows = await getChildrenOf(model, { label: "Agents", id: "AITOOLS.agents", }); @@ -395,7 +395,7 @@ describe(__filename, () => { }); it("excludes the hidden Cursor agent from the installed count", async () => { - const manager = createManager( + const model = createModel( "project", "upToDate", "0.2.9", @@ -406,7 +406,7 @@ describe(__filename, () => { agent("cursor", "Cursor", "0.3.0"), ] ); - const children = await getChildrenOf(manager, { + const children = await getChildrenOf(model, { label: "AI tools", id: "AITOOLS", }); diff --git a/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.ts b/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.ts index f5968470e..0bdee6d26 100644 --- a/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.ts +++ b/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.ts @@ -1,11 +1,8 @@ import {ThemeColor, ThemeIcon, TreeItemCollapsibleState} from "vscode"; import {BaseComponent} from "./BaseComponent"; import {ConfigurationTreeItem} from "./types"; -import { - AiToolsManager, - AiToolsUpdateStatus, - CURSOR_AGENT_ID, -} from "../../aitools/AiToolsManager"; +import {CURSOR_AGENT_ID} from "../../aitools/AiToolsManager"; +import {AiToolsModel, AiToolsUpdateStatus} from "../../aitools/AiToolsModel"; import {HostUtils} from "../../utils"; const TREE_ICON_ID = "AITOOLS"; @@ -33,10 +30,10 @@ function isHiddenAgent(agentId: string): boolean { } export class AiToolsComponent extends BaseComponent { - constructor(private readonly aiToolsManager: AiToolsManager) { + constructor(private readonly aiToolsModel: AiToolsModel) { super(); this.disposables.push( - this.aiToolsManager.onDidChange(() => { + this.aiToolsModel.onDidChange(() => { this.onDidChangeEmitter.fire(); }) ); @@ -44,7 +41,7 @@ export class AiToolsComponent extends BaseComponent { private getRoot(): ConfigurationTreeItem[] { const {installLocation, updateStatus, version, detectError} = - this.aiToolsManager.state; + this.aiToolsModel.state; // Detection failed with an unexpected error and we have no cached // location to fall back on -> show a reload affordance rather than @@ -115,7 +112,7 @@ export class AiToolsComponent extends BaseComponent { parent?: ConfigurationTreeItem ): Promise { const {installLocation, version, detectError, agents} = - this.aiToolsManager.state; + this.aiToolsModel.state; // Only the tree root gets the AI tools row. Guarding solely on // `parent === undefined` is important: ConfigurationDataProvider fans // every getChildren(parent) call out to all components and flattens the diff --git a/packages/databricks-vscode/src/ui/configuration-view/ConfigurationDataProvider.ts b/packages/databricks-vscode/src/ui/configuration-view/ConfigurationDataProvider.ts index 6c6ad62e5..6a3df14dc 100644 --- a/packages/databricks-vscode/src/ui/configuration-view/ConfigurationDataProvider.ts +++ b/packages/databricks-vscode/src/ui/configuration-view/ConfigurationDataProvider.ts @@ -57,7 +57,7 @@ export class ConfigurationDataProvider ) { this.components = [ new WorkspaceFolderComponent(this.workspaceFolderManager), - new AiToolsComponent(this.aiToolsManager), + new AiToolsComponent(this.aiToolsManager.model), new BundleTargetComponent(this.configModel), new AuthTypeComponent( this.connectionManager, diff --git a/packages/databricks-vscode/src/utils/hostUtils.test.ts b/packages/databricks-vscode/src/utils/hostUtils.test.ts index 919d8e6ea..07e42ccf0 100644 --- a/packages/databricks-vscode/src/utils/hostUtils.test.ts +++ b/packages/databricks-vscode/src/utils/hostUtils.test.ts @@ -5,35 +5,30 @@ import {isCursor} from "./hostUtils"; describe(__filename, () => { let originalAppName: PropertyDescriptor | undefined; - function stubAppName(value: string) { - Object.defineProperty(env, "appName", { + function stubUriScheme(value: string) { + Object.defineProperty(env, "uriScheme", { value, configurable: true, }); } beforeEach(() => { - originalAppName = Object.getOwnPropertyDescriptor(env, "appName"); + originalAppName = Object.getOwnPropertyDescriptor(env, "uriScheme"); }); afterEach(() => { if (originalAppName) { - Object.defineProperty(env, "appName", originalAppName); + Object.defineProperty(env, "uriScheme", originalAppName); } }); - it("is true when the app name is Cursor", () => { - stubAppName("Cursor"); + it("is true for Cursor", () => { + stubUriScheme("cursor"); assert.strictEqual(isCursor(), true); }); - it("matches case-insensitively and as a substring", () => { - stubAppName("cursor nightly"); - assert.strictEqual(isCursor(), true); - }); - - it("is false for plain VS Code", () => { - stubAppName("Visual Studio Code"); + it("is false for VS Code", () => { + stubUriScheme("vscode"); assert.strictEqual(isCursor(), false); }); }); diff --git a/packages/databricks-vscode/src/utils/hostUtils.ts b/packages/databricks-vscode/src/utils/hostUtils.ts index 5c352a63b..32edc2a63 100644 --- a/packages/databricks-vscode/src/utils/hostUtils.ts +++ b/packages/databricks-vscode/src/utils/hostUtils.ts @@ -1,10 +1,9 @@ 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. + * Cursor identifies itself via env.uriScheme === "cursor", + * everything else (VS Code, Insiders) uses vscode. */ export function isCursor(): boolean { - return /cursor/i.test(env.appName); + return env.uriScheme === "cursor"; } diff --git a/packages/databricks-vscode/src/vscode-objs/CustomWhenContext.ts b/packages/databricks-vscode/src/vscode-objs/CustomWhenContext.ts index 9eec1133c..68e158b5a 100644 --- a/packages/databricks-vscode/src/vscode-objs/CustomWhenContext.ts +++ b/packages/databricks-vscode/src/vscode-objs/CustomWhenContext.ts @@ -124,4 +124,20 @@ export class CustomWhenContext { value ); } + + setAiToolsInstalled(value: boolean) { + commands.executeCommand( + "setContext", + "databricks.context.aitools.installed", + value + ); + } + + setAiToolsShowCursorPlugin(value: boolean) { + commands.executeCommand( + "setContext", + "databricks.context.aitools.showCursorPlugin", + value + ); + } } diff --git a/packages/databricks-vscode/src/vscode-objs/StateStorage.ts b/packages/databricks-vscode/src/vscode-objs/StateStorage.ts index 69d326d95..c07ac155e 100644 --- a/packages/databricks-vscode/src/vscode-objs/StateStorage.ts +++ b/packages/databricks-vscode/src/vscode-objs/StateStorage.ts @@ -114,8 +114,8 @@ const StorageConfigurations = { // dismissal leaves this false so the prompt can reappear on a later // activation. "databricks.aitools.hideInstallPrompt": withType()({ - location: "global", - defaultValue: false, + location: "global", + defaultValue: false, }), }; From 4d601c1df77a24a76652c7c1f3f8b46f89c713f7 Mon Sep 17 00:00:00 2001 From: Russell Clarey Date: Tue, 4 Aug 2026 13:28:18 +0200 Subject: [PATCH 6/6] remove dev testing code --- packages/databricks-vscode/package.json | 10 -- packages/databricks-vscode/src/extension.ts | 24 ---- .../databricks-vscode/src/test/runTest.ts | 3 - packages/databricks-vscode/src/test/suite.ts | 2 - .../src/vscode-objs/StateResetCommand.test.ts | 130 ------------------ .../src/vscode-objs/StateResetCommand.ts | 56 -------- .../src/vscode-objs/StateStorage.test.ts | 55 -------- .../src/vscode-objs/StateStorage.ts | 25 +--- 8 files changed, 1 insertion(+), 304 deletions(-) delete mode 100644 packages/databricks-vscode/src/vscode-objs/StateResetCommand.test.ts delete mode 100644 packages/databricks-vscode/src/vscode-objs/StateResetCommand.ts diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index 47267cd6c..643985b60 100644 --- a/packages/databricks-vscode/package.json +++ b/packages/databricks-vscode/package.json @@ -140,12 +140,6 @@ "category": "Databricks", "icon": "$(cloud-download)" }, - { - "command": "databricks.developer.resetState", - "title": "Reset State (Developer)", - "category": "Databricks", - "icon": "$(debug-restart)" - }, { "command": "databricks.cluster.filterByAll", "title": "All", @@ -1222,10 +1216,6 @@ } ], "commandPalette": [ - { - "command": "databricks.developer.resetState", - "when": "databricks.context.development" - }, { "command": "databricks.aitools.addCursorPlugin", "when": "databricks.context.aitools.showCursorPlugin && !databricks.context.remoteMode" diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index 84f722a45..1e2ee23c8 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -3,7 +3,6 @@ import { debug, env, ExtensionContext, - ExtensionMode, extensions, window, workspace, @@ -43,7 +42,6 @@ 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, @@ -378,28 +376,6 @@ export async function activate( // Non-blocking so it doesn't delay activation. aiToolsCommands.initializeCommand()(); - // 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 diff --git a/packages/databricks-vscode/src/test/runTest.ts b/packages/databricks-vscode/src/test/runTest.ts index d88a6ae24..6102316e6 100644 --- a/packages/databricks-vscode/src/test/runTest.ts +++ b/packages/databricks-vscode/src/test/runTest.ts @@ -50,9 +50,6 @@ 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 e5ba596a0..2d0a05894 100644 --- a/packages/databricks-vscode/src/test/suite.ts +++ b/packages/databricks-vscode/src/test/suite.ts @@ -7,8 +7,6 @@ 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/vscode-objs/StateResetCommand.test.ts b/packages/databricks-vscode/src/vscode-objs/StateResetCommand.test.ts deleted file mode 100644 index cb55078a7..000000000 --- a/packages/databricks-vscode/src/vscode-objs/StateResetCommand.test.ts +++ /dev/null @@ -1,130 +0,0 @@ -import assert from "assert"; -import {QuickPickItem, commands, window} from "vscode"; -import {anything, capture, instance, mock, verify, when} from "ts-mockito"; -import {StateStorage, StorageKey} from "./StateStorage"; -import {StateResetCommand} from "./StateResetCommand"; - -describe(__filename, () => { - let mockStorage: StateStorage; - let originalShowQuickPick: typeof window.showQuickPick; - let originalShowInfo: typeof window.showInformationMessage; - let originalExecuteCommand: typeof commands.executeCommand; - let executed: string[]; - - // The picker items offered to the user, captured from the last - // showQuickPick call so tests can inspect ordering/labels. - let offeredItems: readonly QuickPickItem[]; - // Behavior applied to the (single) showQuickPick call. - let pickBehavior: ( - items: readonly QuickPickItem[] - ) => QuickPickItem[] | undefined; - // Behavior applied to the reload prompt. - let reloadChoice: string | undefined; - - beforeEach(() => { - mockStorage = mock(StateStorage); - when(mockStorage.storageKeys).thenReturn([ - { - key: "databricks.bundle.target" as StorageKey, - location: "workspace", - }, - { - key: "databricks.aitools.hideInstallPrompt" as StorageKey, - location: "global", - }, - ]); - when(mockStorage.reset(anything())).thenResolve(); - - offeredItems = []; - pickBehavior = () => undefined; - reloadChoice = undefined; - executed = []; - - originalShowQuickPick = window.showQuickPick; - (window as any).showQuickPick = async (items: QuickPickItem[]) => { - offeredItems = items; - return pickBehavior(items); - }; - originalShowInfo = window.showInformationMessage; - (window as any).showInformationMessage = async () => reloadChoice; - originalExecuteCommand = commands.executeCommand; - (commands as any).executeCommand = async (command: string) => { - executed.push(command); - }; - }); - - afterEach(() => { - (window as any).showQuickPick = originalShowQuickPick; - (window as any).showInformationMessage = originalShowInfo; - (commands as any).executeCommand = originalExecuteCommand; - }); - - function createCommand() { - return new StateResetCommand(instance(mockStorage)); - } - - it("offers every storage key, sorted by label, with its location", async () => { - await createCommand().resetCommand()(); - - assert.deepStrictEqual( - offeredItems.map((i) => i.label), - ["databricks.aitools.hideInstallPrompt", "databricks.bundle.target"] - ); - assert.deepStrictEqual( - offeredItems.map((i) => i.description), - ["global", "workspace"] - ); - }); - - it("resets each selected key and offers a window reload", async () => { - pickBehavior = (items) => - items.filter( - (i) => i.label === "databricks.aitools.hideInstallPrompt" - ); - reloadChoice = "Reload Window"; - - await createCommand().resetCommand()(); - - // `reset` is generic; cast for ts-mockito's capture overload. - const [resetKey] = capture(mockStorage.reset as any).last(); - assert.strictEqual(resetKey, "databricks.aitools.hideInstallPrompt"); - verify(mockStorage.reset(anything())).once(); - assert.deepStrictEqual(executed, ["workbench.action.reloadWindow"]); - }); - - it("resets all selected keys", async () => { - pickBehavior = (items) => [...items]; - - await createCommand().resetCommand()(); - - verify(mockStorage.reset(anything())).twice(); - }); - - it("does not reload when the reload prompt is dismissed", async () => { - pickBehavior = (items) => [items[0]]; - reloadChoice = undefined; - - await createCommand().resetCommand()(); - - verify(mockStorage.reset(anything())).once(); - assert.deepStrictEqual(executed, []); - }); - - it("does nothing when the picker is dismissed", async () => { - pickBehavior = () => undefined; - - await createCommand().resetCommand()(); - - verify(mockStorage.reset(anything())).never(); - assert.deepStrictEqual(executed, []); - }); - - it("does nothing when the selection is empty", async () => { - pickBehavior = () => []; - - await createCommand().resetCommand()(); - - verify(mockStorage.reset(anything())).never(); - assert.deepStrictEqual(executed, []); - }); -}); diff --git a/packages/databricks-vscode/src/vscode-objs/StateResetCommand.ts b/packages/databricks-vscode/src/vscode-objs/StateResetCommand.ts deleted file mode 100644 index bebfba345..000000000 --- a/packages/databricks-vscode/src/vscode-objs/StateResetCommand.ts +++ /dev/null @@ -1,56 +0,0 @@ -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, - })) - .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 index 1276a5200..2fbfc88f7 100644 --- a/packages/databricks-vscode/src/vscode-objs/StateStorage.test.ts +++ b/packages/databricks-vscode/src/vscode-objs/StateStorage.test.ts @@ -29,61 +29,6 @@ function createStorage() { 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 hideInstallPrompt = keys.find( - (k) => k.key === "databricks.aitools.hideInstallPrompt" - ); - assert.strictEqual(hideInstallPrompt?.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.hideInstallPrompt", true); - assert.strictEqual( - storage.get("databricks.aitools.hideInstallPrompt"), - true - ); - - await storage.reset("databricks.aitools.hideInstallPrompt"); - - // Raw entry removed, and get falls back to the configured default. - assert.strictEqual( - globalState.get("databricks.aitools.hideInstallPrompt"), - undefined - ); - assert.strictEqual( - storage.get("databricks.aitools.hideInstallPrompt"), - 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 - ); - }); -}); - describe("StateStorage python-setup setup state", () => { it("round-trips the persisted setup state", async () => { const {storage} = createStorage(); diff --git a/packages/databricks-vscode/src/vscode-objs/StateStorage.ts b/packages/databricks-vscode/src/vscode-objs/StateStorage.ts index c07ac155e..4493dc100 100644 --- a/packages/databricks-vscode/src/vscode-objs/StateStorage.ts +++ b/packages/databricks-vscode/src/vscode-objs/StateStorage.ts @@ -119,8 +119,7 @@ const StorageConfigurations = { }), }; -export type StorageKey = keyof typeof StorageConfigurations; -type Keys = StorageKey; +type Keys = keyof typeof StorageConfigurations; type ValueType = (typeof StorageConfigurations)[K]["_type"]; type GetterReturnType> = D extends {getter: infer G} ? G extends (...args: any[]) => any @@ -220,26 +219,4 @@ 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(); - } }