diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index 9a16028fa..7f25b8a22 100644 --- a/packages/databricks-vscode/package.json +++ b/packages/databricks-vscode/package.json @@ -98,6 +98,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.aitools.installAgent", + "title": "Install AI tools for this agent", + "category": "Databricks", + "icon": "$(cloud-download)" + }, { "command": "databricks.cluster.filterByAll", "title": "All", @@ -441,6 +483,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", @@ -1009,6 +1056,21 @@ "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.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", @@ -1161,6 +1223,34 @@ } ], "commandPalette": [ + { + "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.aitools.installAgent", + "when": "false" + }, { "command": "databricks.environment.setupPythonEnv", "when": "false" 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..4d2621b0d --- /dev/null +++ b/packages/databricks-vscode/src/aitools/AiToolsCommands.test.ts @@ -0,0 +1,493 @@ +/* eslint-disable @typescript-eslint/naming-convention */ + +import assert from "assert"; +import { + CancellationToken, + MessageOptions, + Progress, + ProgressOptions, + QuickPick, + QuickPickItem, +} from "vscode"; +import {anything, capture, instance, mock, verify, when} from "ts-mockito"; +import {ProcessError} from "../cli/CliWrapper"; +import {AiToolsManager, CURSOR_AGENT_ID} from "./AiToolsManager"; +import {AiToolsAgentStatus} from "./AiToolsModel"; +import {AiToolsCommands, AiToolsPrompter} 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, + }; +} + +// 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" + > = []; + // 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(() => { + // Default to plain VS Code; Cursor-specific tests opt in via + // stubIsCursor(true). + originalIsCursor = HostUtils.isCursor; + stubIsCursor(false); + }); + + afterEach(() => { + (HostUtils as any).isCursor = originalIsCursor; + }); + + 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), + 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 commands.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 () => { + const {commands, mockManager, behaviors, quickPicks} = setup(); + 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 commands.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); + const {commands, mockManager, behaviors, quickPicks} = setup(); + 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 commands.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); + 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 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), + ]); + 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 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 commands.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 () => { + const {commands, mockManager, behaviors} = setup(); + when(mockManager.listAgents("global")).thenResolve([ + agent("claude-code", "Claude Code", true), + ]); + behaviors.push(selectScope("global")); + behaviors.push(() => "dismiss"); + + 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 commands.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 () => { + 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 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 commands.installAgentCommand()({ + id: "AITOOLS.agent.codex", + }); + + const [agentId] = capture(mockManager.installAgent).last(); + assert.strictEqual(agentId, "codex"); + }); + + it("ignores a node without an agent id", async () => { + const {commands, mockManager} = setup(); + await commands.installAgentCommand()({id: "AITOOLS"}); + await commands.installAgentCommand()(undefined); + + verify(mockManager.installAgent(anything(), anything())).never(); + }); + }); + + describe("addCursorPluginCommand", () => { + it("prompts the plugin with the 'pluginButton' source", async () => { + const {commands, mockManager} = setup(); + when(mockManager.addCursorPlugin(anything())).thenResolve(); + + await commands.addCursorPluginCommand()(); + + const [source] = capture(mockManager.addCursorPlugin).last(); + assert.strictEqual(source, "pluginButton"); + }); + }); + + 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([]); + prompter.messageResponses.push("Install AI tools"); + behaviors.push(selectScope("global")); + + await commands.initializeCommand()(); + + // 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 () => { + const {commands, mockManager, prompter} = setup(); + when(mockManager.initialize()).thenResolve("promptInstall"); + when(mockManager.optOutOfInstallPrompt()).thenResolve(); + prompter.messageResponses.push("Don't show again"); + + await commands.initializeCommand()(); + + verify(mockManager.optOutOfInstallPrompt()).once(); + verify( + mockManager.install(anything(), anything(), anything()) + ).never(); + }); + + it("does not opt out or install on a plain dismissal", async () => { + const {commands, mockManager, prompter} = setup(); + when(mockManager.initialize()).thenResolve("promptInstall"); + prompter.messageResponses.push(undefined); + + await commands.initializeCommand()(); + + verify(mockManager.optOutOfInstallPrompt()).never(); + verify( + mockManager.install(anything(), anything(), anything()) + ).never(); + }); + + it("applies the update when one is available", async () => { + const {commands, mockManager} = setup(); + when(mockManager.initialize()).thenResolve("update"); + when(mockManager.update(anything())).thenResolve(); + + 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 commands.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 {commands, mockManager} = setup(); + 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 commands.updateCommand()(); + + assert.strictEqual( + shownPrefix, + "Failed to update Databricks AI tools." + ); + }); + + it("rethrows a non-ProcessError from update", async () => { + const {commands, mockManager} = setup(); + when(mockManager.update(anything())).thenReject( + new Error("unexpected") + ); + + await assert.rejects( + () => commands.updateCommand()(), + /unexpected/ + ); + }); + }); +}); diff --git a/packages/databricks-vscode/src/aitools/AiToolsCommands.ts b/packages/databricks-vscode/src/aitools/AiToolsCommands.ts new file mode 100644 index 000000000..246376ca4 --- /dev/null +++ b/packages/databricks-vscode/src/aitools/AiToolsCommands.ts @@ -0,0 +1,343 @@ +import { + ProgressLocation, + type CancellationToken, + type Disposable, + type QuickPickItem, + type window, +} from "vscode"; +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"; + +interface ScopeQuickPickItem extends QuickPickItem { + scope: AiToolsScope; +} + +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, + private readonly prompter: AiToolsPrompter + ) {} + + dispose() { + 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 this.prompter.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 this.prompter.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") => { + 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 + * 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; + const quickPick = this.prompter.createQuickPick(); + quickPick.title = "Install Databricks AI tools"; + quickPick.placeholder = "Choose where to install the AI tools"; + quickPick.items = [ + { + label: "$(globe) Global", + detail: "Available to you across all projects", + scope: "global", + }, + { + label: "$(folder) Project", + detail: hasFolder + ? "Checked into the repo, shared with everyone on the project" + : "Open a folder to install AI tools into a project", + // 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", + }, + ]; + + 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(); + }); + } + + /** + * 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 = this.prompter.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.resolveInstalled(); + }; + } + + /** + * Re-run the activation flow. Used by the error row to recover after a + * transient detection failure. + */ + reloadCommand() { + 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.runUpdate(); + }; + } + + uninstallCommand() { + return async () => { + const location = this.aiToolsManager.model.installLocation; + if (location === undefined) { + return; + } + const confirm = await this.prompter.showWarningMessage( + `Uninstall Databricks AI tools (${location})?`, + {modal: true}, + "Uninstall" + ); + if (confirm !== "Uninstall") { + return; + } + await this.withProgress( + "Uninstalling Databricks AI tools", + "Failed to uninstall Databricks AI tools.", + (token) => this.aiToolsManager.uninstall(token) + ); + }; + } + + addCursorPluginCommand() { + return async () => { + 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; + } + 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 new file mode 100644 index 000000000..cf153c5d3 --- /dev/null +++ b/packages/databricks-vscode/src/aitools/AiToolsManager.test.ts @@ -0,0 +1,914 @@ +/* eslint-disable @typescript-eslint/naming-convention */ + +import assert from "assert"; +import {anything, capture, instance, mock, verify, when} from "ts-mockito"; +import {commands, Uri} from "vscode"; +import path from "path"; +import { + AiToolsAgent, + AiToolsListResult, + CliWrapper, + ProcessError, +} 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, + 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<{ + name: string; + latest_version: string; + installed: Record; + }>, + agents: AiToolsAgent[] = [] +): AiToolsListResult { + return { + release: "0.2.9", + skills: skills.map((s) => ({ + experimental: false, + ...s, + })), + agents, + }; +} + +// 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 originalIsCursor: typeof HostUtils.isCursor; + + beforeEach(() => { + // Default to plain VS Code; Cursor-specific tests opt in via + // stubIsCursor(true). + originalIsCursor = HostUtils.isCursor; + stubIsCursor(false); + }); + + afterEach(() => { + (HostUtils as any).isCursor = originalIsCursor; + }); + + it("detects no install when no state file exists", async () => { + const {manager, stubStateStorage} = setup(); + const location = await manager.detectInstall(); + assert.strictEqual(location, undefined); + assert.strictEqual(manager.isInstalled, false); + assert.strictEqual( + stubStateStorage.get("databricks.aitools.installLocation"), + undefined + ); + }); + + it("detects a project install", async () => { + const {manager, stubStateStorage} = setup({project: loadSuccess}); + const location = await manager.detectInstall(); + assert.strictEqual(location, "project"); + assert.strictEqual(manager.isInstalled, true); + assert.strictEqual( + stubStateStorage.get("databricks.aitools.installLocation"), + "project" + ); + }); + + it("detects a global install when only the home state file exists", async () => { + const {manager, stubStateStorage} = setup({global: loadSuccess}); + const location = await manager.detectInstall(); + assert.strictEqual(location, "global"); + assert.strictEqual( + stubStateStorage.get("databricks.aitools.installLocation"), + "global" + ); + }); + + it("prefers project over global when both exist", async () => { + 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. + const loaders: ScopeLoaders = {project: loadSuccess}; + const {manager, stubStateStorage} = setup(loaders); + assert.strictEqual(await manager.detectInstall(), "project"); + assert.strictEqual(manager.model.state.detectError ?? false, false); + + // 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.model.state.installLocation, "project"); + assert.strictEqual(manager.model.state.detectError, true); + assert.strictEqual( + stubStateStorage.get("databricks.aitools.installLocation"), + "project" + ); + }); + + it("clears the detect error flag on a subsequent successful detect", async () => { + const loaders: ScopeLoaders = {project: loadSuccess}; + const {manager} = setup(loaders); + await manager.detectInstall(); + + // Trigger an unexpected read error, then recover. + loaders.project = throwReadError; + await manager.detectInstall(); + assert.strictEqual(manager.model.state.detectError, true); + + // Restore a readable state file; detection should succeed and clear the flag. + loaders.project = loadSuccess; + await manager.detectInstall(); + assert.strictEqual(manager.model.state.detectError, false); + assert.strictEqual(manager.model.state.installLocation, "project"); + }); + + it("reports upToDate when all installed skills match latest", async () => { + const {manager, mockCli} = setup({project: loadSuccess}); + when(mockCli.aitoolsList(anything())).thenResolve( + listResult([ + { + name: "databricks-core", + latest_version: "0.1.0", + installed: {project: "0.1.0"}, + }, + ]) + ); + await manager.detectInstall(); + await manager.resolveInstalled(); + assert.strictEqual(manager.model.state.updateStatus, "upToDate"); + }); + + it("reports updateAvailable when an installed skill is behind latest", async () => { + const {manager, mockCli} = setup({project: loadSuccess}); + 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"}, + }, + ]) + ); + await manager.detectInstall(); + await manager.resolveInstalled(); + assert.strictEqual(manager.model.state.updateStatus, "updateAvailable"); + }); + + it("ignores non-installed skills when computing update status", async () => { + const {manager, mockCli} = setup({project: loadSuccess}); + 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: {}, + }, + ]) + ); + await manager.detectInstall(); + await manager.resolveInstalled(); + assert.strictEqual(manager.model.state.updateStatus, "upToDate"); + }); + + it("reports error when the list command fails", async () => { + const {manager, mockCli} = setup({project: loadSuccess}); + when(mockCli.aitoolsList(anything())).thenReject(new Error("boom")); + await manager.detectInstall(); + await manager.resolveInstalled(); + assert.strictEqual(manager.model.state.updateStatus, "error"); + }); + + it("returns unknown update status when not installed", async () => { + const {manager, mockCli} = setup(); + await manager.detectInstall(); + await manager.resolveInstalled(); + assert.strictEqual(manager.model.state.updateStatus, "unknown"); + verify(mockCli.aitoolsList(anything())).never(); + }); + + it("uninstalls for the detected scope and re-detects", async () => { + const loaders: ScopeLoaders = {project: loadSuccess}; + const {manager, mockCli, stubStateStorage} = setup(loaders); + when( + mockCli.aitoolsUninstall("project", anything(), anything()) + ).thenCall(async () => { + loaders.project = throwNotFound; + }); + 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( + stubStateStorage.get("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 { + const loaders: ScopeLoaders = {project: loadSuccess}; + const {manager, mockCli} = setup(loaders); + when( + mockCli.aitoolsUninstall("project", anything(), anything()) + ).thenCall(async () => { + loaders.project = throwNotFound; + }); + + 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("does not call the CLI when uninstalling with nothing installed", async () => { + const {manager, mockCli} = setup(); + await manager.detectInstall(); + await manager.uninstall(); + verify( + mockCli.aitoolsUninstall(anything(), anything(), anything()) + ).never(); + }); + + 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 () => { + loaders.global = loadSuccess; + }); + when(mockCli.aitoolsList(anything())).thenResolve( + listResult([ + { + name: "databricks-core", + latest_version: "0.1.0", + installed: {global: "0.1.0"}, + }, + ]) + ); + + await manager.install("global"); + + verify( + mockCli.aitoolsInstall("global", anything(), anything(), anything()) + ).once(); + assert.strictEqual(manager.model.state.installLocation, "global"); + assert.strictEqual(manager.model.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 () => { + const {manager, mockCli, stubTelemetry} = setup({ + project: loadSuccess, + }); + 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"}, + }, + ]) + ); + + 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] = 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] = stubTelemetry.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 () => { + const {manager, mockCli, stubTelemetry} = setup({ + project: loadSuccess, + }); + + 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] = stubTelemetry.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 () => { + const {manager, mockCli} = setup({project: loadSuccess}); + 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"}, + }, + ]) + ); + + 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, stubTelemetry} = setup(); + + await manager.addCursorPlugin("pluginButton"); + + const [pluginEvent] = stubTelemetry.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, stubTelemetry} = setup(); + + await manager.addCursorPlugin(); + + assert.deepStrictEqual( + 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 () => { + const {manager, mockCli, stubTelemetry} = setup({project: loadSuccess}); + 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"}, + }, + ]) + ); + await manager.detectInstall(); + + await manager.installAgent("codex"); + + 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] = 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, mockCli} = setup(); + await manager.detectInstall(); + + await manager.installAgent("codex"); + + verify( + mockCli.aitoolsInstall( + anything(), + anything(), + anything(), + anything() + ) + ).never(); + }); + + 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. + 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"}, + }, + ]) + ); + await manager.detectInstall(); + + await assert.rejects(() => manager.installAgent("codex"), ProcessError); + + // resolveInstalled ran despite the failure. + verify(mockCli.aitoolsList(anything())).once(); + 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. + loaders.global = loadSuccess; + throw new ProcessError("boom", 1); + }); + when(mockCli.aitoolsList(anything())).thenResolve( + listResult([ + { + name: "databricks-core", + latest_version: "0.1.0", + installed: {global: "0.1.0"}, + }, + ]) + ); + + await assert.rejects( + () => manager.install("global", "sidePane", ["codex"]), + ProcessError + ); + + // detectInstall + resolveInstalled ran despite the failure, so the row + // reflects the tools that actually installed. + 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 () => { + const {manager, mockCli} = setup({project: loadSuccess}); + 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"}, + }, + ]) + ); + await manager.detectInstall(); + + await manager.update(); + + assert.strictEqual(manager.model.state.updateStatus, "upToDate"); + verify(mockCli.aitoolsList(anything())).once(); + }); + + it("still refreshes update status when the update command fails", async () => { + const {manager, mockCli} = setup({project: loadSuccess}); + 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"}, + }, + ]) + ); + await manager.detectInstall(); + + await assert.rejects(() => manager.update(), ProcessError); + + // The finally block reconciles state even though the update errored. + assert.strictEqual(manager.model.state.updateStatus, "updateAvailable"); + verify(mockCli.aitoolsList(anything())).once(); + }); + + it("captures the installed release version from list", async () => { + const {manager, mockCli} = setup({project: loadSuccess}); + 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"}, + }, + ], + agents: [], + }); + await manager.detectInstall(); + await manager.resolveInstalled(); + assert.strictEqual(manager.model.state.version, "0.3.1"); + }); + + it("clears the version when nothing is installed", async () => { + const loaders: ScopeLoaders = {project: loadSuccess}; + const {manager, mockCli} = setup(loaders); + when(mockCli.aitoolsList(anything())).thenResolve( + listResult([ + { + name: "databricks-core", + latest_version: "0.1.0", + installed: {project: "0.1.0"}, + }, + ]) + ); + await manager.detectInstall(); + await manager.resolveInstalled(); + assert.strictEqual(manager.model.state.version, "0.2.9"); + + // Uninstalling / re-detecting with no state file clears the version. + loaders.project = throwNotFound; + await manager.detectInstall(); + assert.strictEqual(manager.model.state.version, undefined); + }); + + describe("no open folder", () => { + // 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(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", + anything(), + anything(), + anything() + ) + ).thenCall(async () => { + loaders.global = loadSuccess; + }); + when(mockCli.aitoolsList(anything())).thenResolve( + listResult([ + { + name: "databricks-core", + latest_version: "0.1.0", + installed: {global: "0.1.0"}, + }, + ]) + ); + + // Must not throw even though no folder is open. + await manager.install("global"); + + verify( + mockCli.aitoolsInstall( + "global", + homeDir, + anything(), + anything() + ) + ).once(); + assert.strictEqual(manager.model.state.installLocation, "global"); + }); + + it("detects a global install with no folder open", async () => { + 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 () => { + const {manager, mockCli} = setup({project: loadSuccess}); + when(mockCli.aitoolsList(anything())).thenResolve( + listResult([ + { + name: "databricks-core", + latest_version: "0.1.0", + installed: {project: "0.0.1"}, + }, + ]) + ); + + const action = await manager.initialize(); + + // initialize resolves status but leaves applying the update to the + // caller (AiToolsCommands). + assert.strictEqual(action, "update"); + assert.strictEqual( + manager.model.state.updateStatus, + "updateAvailable" + ); + verify( + mockCli.aitoolsUpdate(anything(), anything(), anything()) + ).never(); + }); + + it("reports 'none' when installed and up to date", async () => { + const {manager, mockCli} = setup({project: loadSuccess}); + when(mockCli.aitoolsList(anything())).thenResolve( + listResult([ + { + name: "databricks-core", + latest_version: "0.1.0", + installed: {project: "0.1.0"}, + }, + ]) + ); + + assert.strictEqual(await manager.initialize(), "none"); + }); + + it("reports 'promptInstall' when not installed", async () => { + const {manager} = setup(); + + assert.strictEqual(await manager.initialize(), "promptInstall"); + }); + + it("reports 'none' when not installed but opted out", async () => { + 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, stubStateStorage} = setup(); + assert.strictEqual(manager.shouldPromptInstall, true); + + await manager.optOutOfInstallPrompt(); + + assert.strictEqual( + 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 new file mode 100644 index 000000000..b3bf8cd66 --- /dev/null +++ b/packages/databricks-vscode/src/aitools/AiToolsManager.ts @@ -0,0 +1,556 @@ +import path from "path"; +import {CancellationToken, commands, Disposable} from "vscode"; +import {logging} from "@databricks/sdk-experimental"; +import { + AiToolsAgent, + AiToolsScope, + AiToolsSkill, + CliWrapper, +} 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, + AiToolsInstallSource, + Events, +} 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"; + +/** + * 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", + "aitools", + "skills", + ".state.json" +); + +/** + * 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"; + +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, + })); +} + +/** + * 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[] = []; + + public readonly model: AiToolsModel; + + constructor( + private readonly cli: CliWrapper, + private readonly stateStorage: StateStorage, + private readonly workspaceFolderManager: WorkspaceFolderManager, + private readonly customWhenContext: CustomWhenContext, + private readonly telemetry: Telemetry, + private readonly loadStateFile: StateFileLoader + ) { + this.model = new AiToolsModel( + this.stateStorage.get("databricks.aitools.installLocation") + ); + this.refreshCursorPluginContext(); + this.refreshInstalledContext(); + } + + get isInstalled(): boolean { + return this.model.isInstalled; + } + + /** + * 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(source?: AiToolsCursorPluginSource): Promise { + const recordEvent = this.telemetry.start( + Events.AITOOLS_CURSOR_PLUGIN_PROMPT + ); + try { + await commands.executeCommand( + "workbench.action.openMarketplaceEditor", + { + pluginId: CURSOR_PLUGIN_ID, + openInstallModal: true, + 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 + ); + } + } + + /** + * 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() { + this.customWhenContext.setAiToolsShowCursorPlugin(HostUtils.isCursor()); + } + + /** + * 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() { + this.customWhenContext.setAiToolsInstalled(this.isInstalled); + } + + dispose() { + this.disposables.forEach((d) => d.dispose()); + this.model.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 this.loadStateFile(this.stateFilePath(scope)); + return true; + } catch (e: unknown) { + if (e instanceof Error && "code" in e && 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 AiToolsModel.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"; + } + } 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.model.update({detectError: true}); + this.refreshInstalledContext(); + return this.model.installLocation; + } + + await this.stateStorage.set( + "databricks.aitools.installLocation", + location + ); + // 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(); + return location; + } + + /** + * Entry point run on activation (and by the error-row retry). Detects the + * install state and then: + * - 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. + * + * 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 { + try { + const location = await this.detectInstall(); + if (location === undefined) { + return this.shouldPromptInstall ? "promptInstall" : "none"; + } + await this.resolveInstalled(); + return this.model.state.updateStatus === "updateAvailable" + ? "update" + : "none"; + } catch (e) { + logging.NamedLogger.getOrCreate(Loggers.Extension).error( + "Failed to initialize Databricks AI tools", + e + ); + return "none"; + } + } + + /** + * 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. + */ + get shouldPromptInstall(): boolean { + return !this.stateStorage.get("databricks.aitools.hideInstallPrompt"); + } + + /** + * 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 + ); + } + + /** + * 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 []; + } + } + + /** + * 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 resolveInstalled(): Promise { + const scope = this.model.installLocation; + if (scope === undefined) { + this.model.update({updateStatus: "unknown"}); + return; + } + + this.model.update({updateStatus: "checking"}); + + try { + const result = await this.cli.aitoolsList(this.cwdForScope(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.model.update({updateStatus: "error"}); + } + } + + /** + * 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 + * 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, + agents?: string[], + token?: CancellationToken + ): Promise { + 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); + try { + await this.cli.aitoolsInstall( + scope, + this.cwdForScope(scope), + token, + cliAgents + ); + recordEvent({ + success: true, + scope, + source, + agents: cliAgents, + cursorPlugin, + }); + } catch (e) { + recordEvent({ + success: false, + scope, + source, + agents: cliAgents, + cursorPlugin, + }); + 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(); + } + } + + /** + * 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, + token?: CancellationToken + ): Promise { + const scope = this.model.installLocation; + if (scope === undefined) { + return; + } + const recordEvent = this.telemetry.start(Events.AITOOLS_INSTALL); + try { + await 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], + }); + throw e; + } finally { + // Reconcile the row even on failure: the install may have partially + // succeeded. + await this.resolveInstalled(); + } + } + + /** + * 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(token?: CancellationToken): Promise { + const scope = this.model.installLocation; + if (scope === undefined) { + return; + } + const recordEvent = this.telemetry.start(Events.AITOOLS_UNINSTALL); + try { + await this.cli.aitoolsUninstall( + scope, + this.cwdForScope(scope), + token + ); + recordEvent({success: true, scope}); + } catch (e) { + recordEvent({success: false, scope}); + throw e; + } finally { + await this.detectInstall(); + } + } + + /** + * 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(token?: CancellationToken): Promise { + const scope = this.model.installLocation; + if (scope === undefined) { + return; + } + const recordEvent = this.telemetry.start(Events.AITOOLS_UPDATE); + this.model.update({updateStatus: "updating"}); + try { + await this.cli.aitoolsUpdate(scope, this.cwdForScope(scope), token); + recordEvent({success: true, scope}); + } catch (e) { + recordEvent({success: false, scope}); + 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 + // "Update available" state once the tools are up to date. + await this.resolveInstalled(); + } + } +} 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/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 226bd0b08..a4fb670dc 100644 --- a/packages/databricks-vscode/src/cli/CliWrapper.test.ts +++ b/packages/databricks-vscode/src/cli/CliWrapper.test.ts @@ -1,12 +1,18 @@ 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} from "node:fs/promises"; +import {writeFile, readFile, mkdtemp, rm} from "node:fs/promises"; import {when, spy, reset, instance, mock} from "ts-mockito"; -import {CliWrapper, getSshConnectCommand, waitForProcess} from "./CliWrapper"; +import { + cancellableExecFile, + CliWrapper, + ProcessError, + getSshConnectCommand, + waitForProcess, +} from "./CliWrapper"; import path from "node:path"; import os from "node:os"; import crypto from "node:crypto"; @@ -58,6 +64,38 @@ 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"); + + // 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}); + } + }); + it("should resolve the platform-specific CLI binary name", () => { const cli = createCliWrapper(); const originalPlatform = process.platform; @@ -349,6 +387,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(); @@ -373,3 +434,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/cli/CliWrapper.ts b/packages/databricks-vscode/src/cli/CliWrapper.ts index 77df0e7c9..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,8 +17,10 @@ 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}; import {removeUndefinedKeys} from "../utils/envVarGenerators"; import {quote} from "shell-quote"; import {BundleVariableModel} from "../bundle/models/BundleVariableModel"; @@ -46,11 +47,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; @@ -59,10 +73,16 @@ 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)` 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()}; } @@ -70,7 +90,8 @@ export const execFile = async ( file: string, args: string[], options: Omit = {}, - cancellationToken?: CancellationToken + cancellationToken?: CancellationToken, + execOptions: ExecFileOptions = {} ): Promise<{ stdout: string; stderr: string; @@ -85,38 +106,17 @@ export const execFile = async ( cmd, escapedArgs, escapedOptions, - cancellationToken + cancellationToken, + execOptions ); }; -export interface Command { - command: string; - args: string[]; -} - -export interface ConfigEntry { - name: string; - host?: URL; - accountId?: string; - workspaceId?: string; - cloud: Cloud; - authType: string; - valid: boolean; -} - -export type SyncType = "full" | "incremental"; - -export type SshConnectCompute = - | {type: "serverless"; accelerator?: string} - | {type: "cluster"; clusterId: string}; - /** * Constructs the `databricks ssh connect` command args for opening a remote * 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. @@ -124,7 +124,7 @@ export type SshConnectCompute = 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. @@ -136,6 +136,64 @@ export function getSshConnectCommand(opts: {compute: SshConnectCompute}): { } return {args}; } + +export interface Command { + command: string; + args: string[]; +} + +export interface ConfigEntry { + name: string; + host?: URL; + accountId?: string; + workspaceId?: string; + cloud: Cloud; + authType: string; + valid: boolean; +} + +export type SyncType = "full" | "incremental"; + +export type SshConnectCompute = + | {type: "serverless"; accelerator?: string} + | {type: "cluster"; clusterId: string}; + +/** 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>; +} + +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( message: string, @@ -144,7 +202,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( @@ -172,7 +244,7 @@ export class ProcessError extends Error { ) .then((choice) => { if (choice === "Show Logs") { - commands.executeCommand("databricks.bundle.showLogs"); + commands.executeCommand(logsCommand); } }); } @@ -500,6 +572,133 @@ 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` 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( + 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, + 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 092e818db..5b24aa322 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"; @@ -15,6 +16,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"; @@ -26,6 +29,7 @@ import {logging} from "@databricks/sdk-experimental"; import {workspaceConfigs} from "./vscode-objs/WorkspaceConfigs"; import { FileUtils, + HostUtils, PackageJsonUtils, TerraformUtils, UrlUtils, @@ -261,6 +265,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); @@ -300,6 +305,11 @@ export async function activate( "databricks.bundle.showLogs", () => loggerManager.showOutputChannel("Databricks Bundle Logs"), loggerManager + ), + telemetry.registerCommand( + "databricks.logs.show", + () => loggerManager.showOutputChannel("Databricks Logs"), + loggerManager ) ); @@ -313,6 +323,67 @@ 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, + customWhenContext, + telemetry, + (path) => readFile(path) + ); + const aiToolsCommands = new AiToolsCommands(aiToolsManager, window); + 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 + ), + telemetry.registerCommand( + "databricks.aitools.installAgent", + aiToolsCommands.installAgentCommand(), + aiToolsCommands + ) + ); + // 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. + aiToolsCommands.initializeCommand()(); + if ( workspace.workspaceFolders === undefined || workspace.workspaceFolders?.length === 0 @@ -323,7 +394,8 @@ export async function activate( async () => { const bundleInitWizard = new BundleInitWizard( cli, - telemetry + telemetry, + aiToolsManager ); await bundleInitWizard.initNewProject(); } @@ -344,15 +416,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 - ); - // Utils. Registered before the remote-mode branch below returns so that // views shared between the normal and remote flows - such as the docs view, // which invokes "databricks.utils.openExternal" - work in both. @@ -662,7 +730,8 @@ export async function activate( configModel, bundleFileSet, workspaceFolderManager, - telemetry + telemetry, + aiToolsManager ); context.subscriptions.push( bundleProjectManager, @@ -1007,6 +1076,7 @@ export async function activate( cli, featureManager, workspaceFolderManager, + aiToolsManager, pythonSetupEnvironment ); const configurationView = window.createTreeView("configurationView", { diff --git a/packages/databricks-vscode/src/telemetry/constants.ts b/packages/databricks-vscode/src/telemetry/constants.ts index 9ed8b89f4..f1e7617f9 100644 --- a/packages/databricks-vscode/src/telemetry/constants.ts +++ b/packages/databricks-vscode/src/telemetry/constants.ts @@ -25,6 +25,10 @@ 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", + AITOOLS_CURSOR_PLUGIN_PROMPT = "aitoolsCursorPluginPrompt", } /* eslint-enable @typescript-eslint/naming-convention */ @@ -43,6 +47,21 @@ 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 init modal prompt + * or the manual affordance in the configuration side pane. + */ +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 @@ -166,9 +185,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 +250,78 @@ export class EventTypes { comment: "The resource type", }, }; + [Events.AITOOLS_INSTALL]: EventType< + { + success: boolean; + scope: AiToolsScope; + source?: AiToolsInstallSource; + agents?: string[]; + cursorPlugin?: boolean; + } & 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: '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< + { + 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.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 new file mode 100644 index 000000000..300a78173 --- /dev/null +++ b/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.test.ts @@ -0,0 +1,420 @@ +/* eslint-disable @typescript-eslint/naming-convention */ + +import assert from "assert"; +import {ThemeIcon} from "vscode"; +import { + AiToolsAgentStatus, + AiToolsInstallLocation, + AiToolsModel, + AiToolsUpdateStatus, +} from "../../aitools/AiToolsModel"; +import {resolveProviderResult} from "../../test/utils"; +import {AiToolsComponent} from "./AiToolsComponent"; +import {HostUtils} from "../../utils"; + +function createModel( + installLocation: AiToolsInstallLocation, + updateStatus: AiToolsUpdateStatus, + version?: string, + detectError?: boolean, + agents: AiToolsAgentStatus[] = [] +): AiToolsModel { + return { + state: {installLocation, updateStatus, version, detectError, agents}, + onDidChange: () => ({dispose() {}}), + } as unknown as AiToolsModel; +} + +async function getRoot(model: AiToolsModel) { + const component = new AiToolsComponent(model); + const items = await resolveProviderResult(component.getChildren()); + return items ?? []; +} + +async function getChildrenOf( + model: AiToolsModel, + parent: {label?: string; id?: string} +) { + const component = new AiToolsComponent(model); + 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(createModel(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( + createModel(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( + createModel("project", "upToDate", "0.2.9", true) + ); + const [row] = items; + assert.strictEqual(row.label, "AI tools"); + assert.ok(String(row.tooltip).includes("project")); + }); + + it("renders the installed version in the subtext for a project install", async () => { + const items = await getRoot( + createModel("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.tooltip).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(createModel("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(createModel("global", "updateAvailable")); + const [row] = items; + assert.strictEqual( + row.contextValue, + "databricks.configuration.aitools.updateAvailable" + ); + 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); + }); + + it("renders an updating spinner while auto-updating", async () => { + const items = await getRoot(createModel("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(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(createModel("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(createModel("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( + createModel("project", "upToDate") + ); + const children = await resolveProviderResult( + 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( + createModel(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( + createModel(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 model = createModel( + "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(model, { + 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 model = createModel( + "project", + "upToDate", + "0.2.9", + undefined, + [agent("cursor", "Cursor")] + ); + const children = await getChildrenOf(model, { + 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 model = createModel( + "project", + "upToDate", + "0.2.9", + undefined, + [] + ); + const children = await getChildrenOf(model, { + 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 model = createModel( + "project", + "upToDate", + "0.2.9", + undefined, + [ + agent("claude", "Claude Code", "1.2.0"), + agent("cursor", "Cursor"), + ] + ); + const rows = await getChildrenOf(model, { + 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 model = createModel( + "project", + "upToDate", + "0.2.9", + undefined, + [] + ); + const rows = await getChildrenOf(model, { + label: "Agents", + id: "AITOOLS.agents", + }); + assert.deepStrictEqual(rows, []); + }); + + it("marks installed agents with a green check and no install affordance", async () => { + const model = createModel( + "project", + "upToDate", + "0.2.9", + undefined, + [agent("claude", "Claude Code", "1.2.0")] + ); + const [row] = await getChildrenOf(model, { + 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 model = createModel( + "project", + "upToDate", + "0.2.9", + undefined, + [agent("codex", "Codex CLI")] + ); + const [row] = await getChildrenOf(model, { + 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 model = createModel( + "project", + "upToDate", + "0.2.9", + undefined, + [agent("codex", "Codex CLI")] + ); + const [row] = await getChildrenOf(model, { + 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"}, + ]); + }); + + 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 model = createModel( + "project", + "upToDate", + "0.2.9", + undefined, + [ + agent("claude", "Claude Code", "1.2.0"), + agent("cursor", "Cursor", "0.3.0"), + ] + ); + const rows = await getChildrenOf(model, { + 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 model = createModel( + "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(model, { + 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/ui/configuration-view/AiToolsComponent.ts b/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.ts new file mode 100644 index 000000000..0bdee6d26 --- /dev/null +++ b/packages/databricks-vscode/src/ui/configuration-view/AiToolsComponent.ts @@ -0,0 +1,258 @@ +import {ThemeColor, ThemeIcon, TreeItemCollapsibleState} from "vscode"; +import {BaseComponent} from "./BaseComponent"; +import {ConfigurationTreeItem} from "./types"; +import {CURSOR_AGENT_ID} from "../../aitools/AiToolsManager"; +import {AiToolsModel, AiToolsUpdateStatus} from "../../aitools/AiToolsModel"; +import {HostUtils} from "../../utils"; + +const TREE_ICON_ID = "AITOOLS"; + +function getTreeIconId(key: string) { + return `${TREE_ICON_ID}.${key}`; +} + +function robotIcon(color: "blue" | "green") { + return new ThemeIcon("hubot", new ThemeColor(`charts.${color}`)); +} + +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 aiToolsModel: AiToolsModel) { + super(); + this.disposables.push( + this.aiToolsModel.onDidChange(() => { + this.onDidChangeEmitter.fire(); + }) + ); + } + + private getRoot(): ConfigurationTreeItem[] { + const {installLocation, updateStatus, version, detectError} = + 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 + // implying the tools simply aren't installed. + if (installLocation === undefined && detectError) { + return [ + { + label: "AI tools", + id: TREE_ICON_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", + }, + }, + ]; + } + + if (installLocation === undefined) { + return [ + { + label: "Install AI tools", + id: TREE_ICON_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: TREE_ICON_ID, + description: description ?? "", + tooltip: `AI tools installed (${installLocation})`, + contextValue: getContextValue(state), + iconPath: icon, + collapsibleState: TreeItemCollapsibleState.Collapsed, + }, + ]; + + // 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 { + const {installLocation, version, detectError, agents} = + 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 + // 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 []; + } + + 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 []; + } +} + +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 { + icon: robotIcon("green"), + description: versionLabel ?? "Up to date", + state: "upToDate", + }; + case "updateAvailable": + return { + icon: robotIcon("green"), + description: "Update available", + state: "updateAvailable", + }; + case "checking": + 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 1719c750a..6f45b7bff 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"; import {PythonSetupEntry} from "./pythonSetupEntry"; /** @@ -52,10 +54,12 @@ export class ConfigurationDataProvider private readonly cli: CliWrapper, private readonly featureManager: FeatureManager, private readonly workspaceFolderManager: WorkspaceFolderManager, + private readonly aiToolsManager: AiToolsManager, private readonly pythonSetup?: PythonSetupEntry ) { this.components = [ new WorkspaceFolderComponent(this.workspaceFolderManager), + 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 new file mode 100644 index 000000000..07e42ccf0 --- /dev/null +++ b/packages/databricks-vscode/src/utils/hostUtils.test.ts @@ -0,0 +1,34 @@ +import {env} from "vscode"; +import assert from "assert"; +import {isCursor} from "./hostUtils"; + +describe(__filename, () => { + let originalAppName: PropertyDescriptor | undefined; + + function stubUriScheme(value: string) { + Object.defineProperty(env, "uriScheme", { + value, + configurable: true, + }); + } + + beforeEach(() => { + originalAppName = Object.getOwnPropertyDescriptor(env, "uriScheme"); + }); + + afterEach(() => { + if (originalAppName) { + Object.defineProperty(env, "uriScheme", originalAppName); + } + }); + + it("is true for Cursor", () => { + stubUriScheme("cursor"); + assert.strictEqual(isCursor(), true); + }); + + 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 new file mode 100644 index 000000000..32edc2a63 --- /dev/null +++ b/packages/databricks-vscode/src/utils/hostUtils.ts @@ -0,0 +1,9 @@ +import {env} from "vscode"; + +/** + * Cursor identifies itself via env.uriScheme === "cursor", + * everything else (VS Code, Insiders) uses vscode. + */ +export function isCursor(): boolean { + return env.uriScheme === "cursor"; +} 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..68e158b5a 100644 --- a/packages/databricks-vscode/src/vscode-objs/CustomWhenContext.ts +++ b/packages/databricks-vscode/src/vscode-objs/CustomWhenContext.ts @@ -116,4 +116,28 @@ export class CustomWhenContext { value ); } + + setIsCursor(value: boolean) { + commands.executeCommand( + "setContext", + "databricks.context.isCursor", + 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.test.ts b/packages/databricks-vscode/src/vscode-objs/StateStorage.test.ts index eea9efce0..2fbfc88f7 100644 --- a/packages/databricks-vscode/src/vscode-objs/StateStorage.test.ts +++ b/packages/databricks-vscode/src/vscode-objs/StateStorage.test.ts @@ -1,4 +1,4 @@ -import {expect} from "chai"; +import assert from "assert"; import {ExtensionContext, Memento} from "vscode"; import {PythonSetupState, StateStorage} from "./StateStorage"; @@ -19,17 +19,19 @@ function fakeMemento(): Memento { } as Memento; } -function makeStorage(): StateStorage { +function createStorage() { + const globalState = fakeMemento(); + const workspaceState = fakeMemento(); const context = { - workspaceState: fakeMemento(), - globalState: fakeMemento(), + globalState, + workspaceState, } as unknown as ExtensionContext; - return new StateStorage(context); + return {storage: new StateStorage(context), globalState, workspaceState}; } describe("StateStorage python-setup setup state", () => { it("round-trips the persisted setup state", async () => { - const storage = makeStorage(); + const {storage} = createStorage(); const state: PythonSetupState = { envKey: "serverless/serverless-v5", pythonVersion: "3.12", @@ -38,19 +40,22 @@ describe("StateStorage python-setup setup state", () => { await storage.set("databricks.pythonSetup.setupState", state); - expect(storage.get("databricks.pythonSetup.setupState")).to.deep.equal( + assert.deepStrictEqual( + storage.get("databricks.pythonSetup.setupState"), state ); }); it("returns undefined before anything is persisted", () => { - expect(makeStorage().get("databricks.pythonSetup.setupState")).to.equal( + const {storage} = createStorage(); + assert.strictEqual( + storage.get("databricks.pythonSetup.setupState"), undefined ); }); it("clears the state when set to undefined", async () => { - const storage = makeStorage(); + const {storage} = createStorage(); await storage.set("databricks.pythonSetup.setupState", { envKey: "cluster/0101", pythonVersion: "3.11", @@ -59,18 +64,14 @@ describe("StateStorage python-setup setup state", () => { await storage.set("databricks.pythonSetup.setupState", undefined); - expect(storage.get("databricks.pythonSetup.setupState")).to.equal( + assert.strictEqual( + storage.get("databricks.pythonSetup.setupState"), undefined ); }); it("is stored in workspace state (per-project), not global", async () => { - const workspaceState = fakeMemento(); - const globalState = fakeMemento(); - const storage = new StateStorage({ - workspaceState, - globalState, - } as unknown as ExtensionContext); + const {storage, globalState, workspaceState} = createStorage(); await storage.set("databricks.pythonSetup.setupState", { envKey: "serverless/serverless-v5", @@ -78,10 +79,12 @@ describe("StateStorage python-setup setup state", () => { timestamp: "2026-07-27T10:00:00.000Z", }); - expect( - workspaceState.get("databricks.pythonSetup.setupState") - ).to.not.equal(undefined); - expect(globalState.get("databricks.pythonSetup.setupState")).to.equal( + assert.notStrictEqual( + workspaceState.get("databricks.pythonSetup.setupState"), + undefined + ); + assert.strictEqual( + globalState.get("databricks.pythonSetup.setupState"), undefined ); }); diff --git a/packages/databricks-vscode/src/vscode-objs/StateStorage.ts b/packages/databricks-vscode/src/vscode-objs/StateStorage.ts index ac0001682..4493dc100 100644 --- a/packages/databricks-vscode/src/vscode-objs/StateStorage.ts +++ b/packages/databricks-vscode/src/vscode-objs/StateStorage.ts @@ -99,6 +99,24 @@ const StorageConfigurations = { "databricks.pythonSetup.setupState": withType()({ location: "workspace", }), + + // 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 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, + }), }; type Keys = keyof typeof StorageConfigurations; 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(); + }); + }); }); 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