diff --git a/src/cosmosDBShell/CosmosDBShellExtension.ts b/src/cosmosDBShell/CosmosDBShellExtension.ts index 49bd99fb3..6837d8b49 100644 --- a/src/cosmosDBShell/CosmosDBShellExtension.ts +++ b/src/cosmosDBShell/CosmosDBShellExtension.ts @@ -4,8 +4,10 @@ *--------------------------------------------------------------------------------------------*/ /** - * Entry point for Cosmos DB Shell support. Manages command registration, terminal orchestration, - * MCP provider wiring, and language server startup. + * Entry point for Cosmos DB Shell support. Owns the {@link CosmosDBShellExtension} + * activation lifecycle and the two top-level command handlers ({@link launchCosmosDBShell} + * and {@link connectCosmosDBShell}). Heavier subsystems (install flow, MCP provider, + * language server, terminal reuse, version cache) live in sibling modules. */ import { type ContainerDefinition, type DatabaseDefinition } from '@azure/cosmos'; import { @@ -14,44 +16,43 @@ import { type IActionContext, } from '@microsoft/vscode-azext-utils'; import * as l10n from '@vscode/l10n'; -import * as child from 'child_process'; -import * as fs from 'fs'; -import * as http from 'http'; -import * as net from 'net'; import * as vscode from 'vscode'; -import { - LanguageClient, - RevealOutputChannelOn, - TransportKind, - type LanguageClientOptions, - type ServerOptions, -} from 'vscode-languageclient/node'; -import { AuthenticationMethod } from '../cosmosdb/AuthenticationMethod'; -import { type CosmosDBEntraIdCredential, type CosmosDBManagedIdentityCredential } from '../cosmosdb/CosmosDBCredential'; -import { getAccessTokenForVSCode } from '../cosmosdb/utils/azureSessionHelper'; import { ext } from '../extensionVariables'; import { SettingsService } from '../services/SettingsService'; import { type NoSqlContainerResourceItem } from '../tree/nosql/NoSqlContainerResourceItem'; -import { resolveCosmosDBShellCommand } from './cosmosDBShellCommandResolver'; -import { CosmosDBShellMcpHost, getCosmosDBShellMcpEndpoint } from './cosmosDBShellMcpEndpoint'; - -type AuthKind = 'emulator' | 'accountKey' | 'entraId' | 'managedIdentity' | 'none'; - -type ShellTerminalState = { - /** Endpoint the shell process was launched against, or '' for command-palette launches without a node. */ - endpoint: string; - /** Authentication mode used at launch. Determines which env vars (if any) are baked into the process. */ - authKind: AuthKind; - tenantId?: string; - managedIdentityClientId?: string; -}; - -// Track per-terminal launch state so we know which Cosmos DB Shell terminals can be reused -// for a given node. The state describes how the process was *launched* (endpoint + auth -// mode + env vars baked in), not its current in-shell connection status: a user may have -// run `disconnect` inside the shell, which VS Code cannot observe. The reuse path therefore -// always re-issues `connect` before sending further commands. -const terminalStates = new Map(); +import { + COMMAND_LAUNCH_COSMOS_DB_SHELL, + COSMOS_DB_SHELL_TERMINAL_NAME, + DEFAULT_MCP_PORT, + SETTING_MCP_ENABLED, + SETTING_MCP_PORT, + SETTING_SHELL_PATH, +} from './constants'; +import { promptToResolveMissingCosmosDBShell } from './install/installPrompts'; +import { + getCosmosDBShellCredential, + getCosmosDBShellToken, + getEntraIdCredential, + getManagedIdentityCredential, + getNodeAuthKind, +} from './nodeCredentials'; +import { getCosmosDBShellCommand, watchForEarlyExit } from './shellCommand'; +import { getDetectedCosmosDBShellVersion, isCosmosDBShellInstalled } from './shellSupportCache'; +import { + buildInteractiveConnectCommand, + buildTerminalStateForNode, + findReusableTerminalForNode, + terminalStates, +} from './terminalReuse'; + +// Re-exports preserve the existing public surface consumed by ../extension.ts. +export { registerCosmosDBShellLanguageServer } from './languageServer'; +export { registerMcpServer } from './mcpProvider'; +export { + getDetectedCosmosDBShellVersion, + invalidateCosmosDBShellSupportCache, + isCosmosDBShellInstalled, +} from './shellSupportCache'; export class CosmosDBShellExtension implements vscode.Disposable { private terminalChangeListeners: vscode.Disposable[] = []; @@ -71,7 +72,7 @@ export class CosmosDBShellExtension implements vscode.Disposable { 'cosmosDB.cosmosDBShell.activate', (_activateContext: IActionContext) => { const shellInstalled: boolean = isCosmosDBShellInstalled(); - vscode.commands.executeCommand( + void vscode.commands.executeCommand( 'setContext', 'vscodeDatabases.cosmosDBShellSupportEnabled', shellInstalled, @@ -83,7 +84,7 @@ export class CosmosDBShellExtension implements vscode.Disposable { // Watch for terminal open events const openListener = vscode.window.onDidOpenTerminal((terminal) => { // Check if it's a Cosmos DB Shell terminal - if (terminal.name === 'Cosmos DB Shell') { + if (terminal.creationOptions.name === COSMOS_DB_SHELL_TERMINAL_NAME) { this.updateCosmosDBShellTerminalContext(); } }); @@ -91,7 +92,7 @@ export class CosmosDBShellExtension implements vscode.Disposable { // Watch for terminal close events const closeListener = vscode.window.onDidCloseTerminal((terminal) => { // Check if it was a Cosmos DB Shell terminal - if (terminal.name === 'Cosmos DB Shell') { + if (terminal.creationOptions.name === COSMOS_DB_SHELL_TERMINAL_NAME) { this.updateCosmosDBShellTerminalContext(); // Remove tracked launch state for this terminal terminalStates.delete(terminal); @@ -101,7 +102,7 @@ export class CosmosDBShellExtension implements vscode.Disposable { // Store listeners for disposal this.terminalChangeListeners.push(openListener, closeListener); - registerCommandWithTreeNodeUnwrapping('cosmosDB.launchCosmosDBShell', connectCosmosDBShell); + registerCommandWithTreeNodeUnwrapping(COMMAND_LAUNCH_COSMOS_DB_SHELL, connectCosmosDBShell); if (shellInstalled) { ext.outputChannel.appendLine(`Cosmos DB Shell Extension: activated.`); @@ -114,9 +115,9 @@ export class CosmosDBShellExtension implements vscode.Disposable { private updateCosmosDBShellTerminalContext(): void { const hasCosmosDBShellTerminal = vscode.window.terminals.some( - (terminal) => terminal.name === 'Cosmos DB Shell', + (terminal) => terminal.creationOptions.name === COSMOS_DB_SHELL_TERMINAL_NAME, ); - vscode.commands.executeCommand( + void vscode.commands.executeCommand( 'setContext', 'vscodeDatabases.cosmosDBShellTerminalOpen', hasCosmosDBShellTerminal, @@ -127,473 +128,38 @@ export class CosmosDBShellExtension implements vscode.Disposable { } } -function getCosmosDBShellCommand(): string { - const shellPath: string | undefined = SettingsService.getSetting('cosmosDB.shell.path'); - return resolveCosmosDBShellCommand(shellPath); -} - -function isCosmosDBShellPathFound(): boolean { - const shellPath: string | undefined = SettingsService.getSetting('cosmosDB.shell.path'); - if (!shellPath?.trim()) { - return false; - } - - const trimmed = shellPath.trim(); - const unquoted = - (trimmed.startsWith('"') && trimmed.endsWith('"')) || (trimmed.startsWith("'") && trimmed.endsWith("'")) - ? trimmed.slice(1, -1) - : trimmed; - - try { - return fs.existsSync(unquoted) && fs.statSync(unquoted).isFile(); - } catch { - return false; - } -} - -/** - * Watches for the terminal closing shortly after creation (early exit). - * If the process exits quickly, logs the exit code and reason to the output channel. - */ -function watchForEarlyExit(terminal: vscode.Terminal): void { - const startTime = Date.now(); - const listener = vscode.window.onDidCloseTerminal((closedTerminal) => { - if (closedTerminal !== terminal) { - return; - } - - clearTimeout(timeout); - listener.dispose(); - if (Date.now() - startTime < 5000) { - const exitCode = closedTerminal.exitStatus?.code; - const exitReason = closedTerminal.exitStatus?.reason; - ext.outputChannel.error( - `Cosmos DB Shell exited early.${exitCode !== undefined ? ` Exit code: ${exitCode}.` : ''}${exitReason !== undefined ? ` Reason: ${exitReason}.` : ''}`, - ); - } - }); - const timeout = setTimeout(() => { - listener.dispose(); - }, 5000); -} - -/** - * Minimum .NET SDK version required to install the Cosmos DB Shell global tool. - * Bumping this constant also updates the version requested via - * `dotnet.acquireGlobalSDK` and the user-facing prompt copy. - */ -const MIN_DOTNET_SDK_VERSION = '10.0.203'; - -/** - * Channel (`major.minor`) requested from `dotnet.acquireGlobalSDK`. Using a - * channel rather than {@link MIN_DOTNET_SDK_VERSION} lets the .NET Install Tool - * resolve to the latest available patch on that channel while still satisfying - * the floor enforced by {@link hasRequiredDotNetSdk}. - */ -const REQUESTED_DOTNET_SDK_CHANNEL = MIN_DOTNET_SDK_VERSION.split('.').slice(0, 2).join('.'); - -/** - * Compares two .NET SDK version strings (e.g. `10.0.203`, `9.0.100-rc.1`) by - * numeric major / minor / patch components. Any pre-release / build metadata - * suffix (after `-`) is ignored. Returns a negative number when `a < b`, zero - * when equal, and a positive number when `a > b`. - */ -function compareDotNetVersions(a: string, b: string): number { - const parse = (v: string): number[] => - v - .split('-')[0] - .split('.') - .map((part) => Number.parseInt(part, 10) || 0); - const av = parse(a); - const bv = parse(b); - const len = Math.max(av.length, bv.length); - for (let i = 0; i < len; i++) { - const diff = (av[i] ?? 0) - (bv[i] ?? 0); - if (diff !== 0) { - return diff; - } - } - return 0; -} - -/** - * Runs `dotnet --list-sdks` and returns the parsed version strings. The CLI - * output format is ` []` per line. Returns an empty - * array when `dotnet` is not on PATH or the call fails. - */ -function getInstalledDotNetSdkVersions(dotnetPath?: string): string[] { - try { - const output = child.execFileSync(dotnetPath ?? 'dotnet', ['--list-sdks'], { - windowsHide: true, - stdio: 'pipe', - }); - return output - .toString('utf8') - .split(/\r?\n/) - .map((line) => line.trim()) - .filter((line) => line.length > 0) - .map((line) => line.split(/\s+/, 1)[0]); - } catch { - return []; - } -} - -function hasRequiredDotNetSdk(dotnetPath?: string): boolean { - return getInstalledDotNetSdkVersions(dotnetPath).some( - (version) => compareDotNetVersions(version, MIN_DOTNET_SDK_VERSION) >= 0, - ); -} - -/** - * Runs `dotnet tool install --global CosmosDBShell --prerelease` with a progress - * notification, streaming output to the extension output channel. Returns true - * when the process exits with code 0. - */ -async function installCosmosDBShellWithDotNetTool(dotnetPath?: string): Promise { - const result = await callWithTelemetryAndErrorHandling( - 'cosmosDB.cosmosDBShell.install.tool', - async (telemetryContext: IActionContext) => { - telemetryContext.errorHandling.suppressDisplay = true; - telemetryContext.telemetry.properties.dotnetPathProvided = String(!!dotnetPath); - const startedAt = Date.now(); - const outcome = await vscode.window.withProgress( - { - location: vscode.ProgressLocation.Notification, - title: l10n.t('Installing Cosmos DB Shell…'), - cancellable: true, - }, - async (_progress, token) => { - ext.outputChannel.show(true); - const dotnetExe = dotnetPath ?? 'dotnet'; - ext.outputChannel.appendLine(`> ${dotnetExe} tool install --global CosmosDBShell --prerelease`); - - return new Promise<{ success: boolean; exitCode: number | null; cancelled: boolean }>((resolve) => { - let cancelled = false; - const proc = child.spawn( - dotnetExe, - ['tool', 'install', '--global', 'CosmosDBShell', '--prerelease'], - { windowsHide: true, shell: false }, - ); - - token.onCancellationRequested(() => { - cancelled = true; - proc.kill(); - }); - - proc.stdout?.on('data', (data: Buffer) => { - ext.outputChannel.append(data.toString('utf8')); - }); - proc.stderr?.on('data', (data: Buffer) => { - ext.outputChannel.append(data.toString('utf8')); - }); - proc.on('error', (err) => { - ext.outputChannel.appendLine(`Failed to start dotnet: ${err.message}`); - resolve({ success: false, exitCode: null, cancelled }); - }); - proc.on('close', (code) => { - ext.outputChannel.appendLine(`\nProcess exited with code ${code}.`); - resolve({ success: code === 0, exitCode: code, cancelled }); - }); - }); - }, - ); - telemetryContext.telemetry.measurements.durationMs = Date.now() - startedAt; - telemetryContext.telemetry.properties.exitCode = - outcome.exitCode === null ? 'null' : String(outcome.exitCode); - telemetryContext.telemetry.properties.cancelled = String(outcome.cancelled); - telemetryContext.telemetry.properties.outcome = outcome.cancelled - ? 'cancelled' - : outcome.success - ? 'success' - : 'failure'; - return outcome.success; - }, - ); - return result ?? false; -} - -/** - * Fires a `cosmosDB.cosmosDBShell.install.prompt` telemetry event with the - * given prompt identifier and user selection. Used to measure the install - * funnel without depending on localized button labels. - */ -function reportInstallPromptOutcome( - promptKind: - | 'missingShell' - | 'installShell' - | 'installSdk' - | 'pathMisconfigured' - | 'reloadAfterInstall' - | 'installFailure', - selection: string, - extraProperties?: Record, -): void { - void callWithTelemetryAndErrorHandling( - 'cosmosDB.cosmosDBShell.install.prompt', - (telemetryContext: IActionContext) => { - telemetryContext.errorHandling.suppressDisplay = true; - telemetryContext.telemetry.properties.promptKind = promptKind; - telemetryContext.telemetry.properties.selection = selection; - if (extraProperties) { - for (const [k, v] of Object.entries(extraProperties)) { - telemetryContext.telemetry.properties[k] = v; - } - } - }, - ); -} - -/** - * Prompts the user to install Cosmos DB Shell via `dotnet tool install`, and on - * success automatically continues the original launch flow. - */ -async function promptToInstallCosmosDBShell( - context: IActionContext, - node: NoSqlContainerResourceItem | undefined, -): Promise { - const install = l10n.t('Install'); - const settings = l10n.t('Settings'); - const selection = await vscode.window.showInformationMessage( - l10n.t( - 'Cosmos DB Shell is not installed. Install it now using `dotnet tool install --global CosmosDBShell --prerelease`?', - ), - { modal: true }, - install, - settings, - ); - - const outcome = selection === install ? 'install' : selection === settings ? 'settings' : 'cancelled'; - reportInstallPromptOutcome('installShell', outcome); - - if (selection === settings) { - void vscode.commands.executeCommand('workbench.action.openSettings', 'cosmosDB.shell.path'); - return; - } - if (selection !== install) { - return; - } - - await installAndLaunchCosmosDBShell(context, node); -} - -/** - * Runs the `dotnet tool install` for Cosmos DB Shell, then either reloads the window - * (if PATH hasn't picked up the new tool yet) or auto-launches the shell to continue - * the user's original action. Used both by the explicit install prompt and by the - * auto-chain after a fresh .NET SDK acquisition. - */ -async function installAndLaunchCosmosDBShell( - context: IActionContext, - node: NoSqlContainerResourceItem | undefined, - dotnetPath?: string, -): Promise { - const success = await installCosmosDBShellWithDotNetTool(dotnetPath); - if (!success) { - const showOutput = l10n.t('Show Output'); - const failureSelection = await vscode.window.showErrorMessage( - l10n.t('Failed to install Cosmos DB Shell. See the output for details.'), - showOutput, - ); - reportInstallPromptOutcome('installFailure', failureSelection === showOutput ? 'showOutput' : 'dismissed'); - if (failureSelection === showOutput) { - ext.outputChannel.show(true); - } - return; - } - - // On a brand-new install the user's PATH may not yet include `~/.dotnet/tools` - // in the current VS Code session. If we still can't resolve the shell, ask to reload. - if (!isCosmosDBShellInstalled()) { - const reload = l10n.t('Reload Window'); - const reloadSelection = await vscode.window.showInformationMessage( - l10n.t( - 'Cosmos DB Shell was installed, but its location is not yet on PATH for this VS Code window. Reload the window to pick it up.', - ), - reload, - ); - reportInstallPromptOutcome('reloadAfterInstall', reloadSelection === reload ? 'reload' : 'cancelled'); - if (reloadSelection === reload) { - void vscode.commands.executeCommand('workbench.action.reloadWindow'); - } - return; - } - - // Auto-relaunch with the original node so the user lands where they intended. - await launchCosmosDBShell(context, node); -} - -/** - * Awaits the `.NET Install Tool` SDK acquisition command, requesting a global - * install of the latest patch available on {@link REQUESTED_DOTNET_SDK_CHANNEL}. - * Returns the resolved `dotnet` executable path on success, or undefined when - * the acquisition failed or did not return a path. - * - * Uses `dotnet.acquireGlobalSDK` rather than `dotnet.acquireGlobalSDKPublic` - * because the public variant ignores the supplied `version` and prompts the - * user with its own recommended version instead. - */ -async function tryInstallDotNetSdkViaExtension(): Promise { - const result = await callWithTelemetryAndErrorHandling( - 'cosmosDB.cosmosDBShell.install.dotnetSdk', - async (telemetryContext: IActionContext) => { - telemetryContext.errorHandling.suppressDisplay = true; - telemetryContext.telemetry.properties.requestedChannel = REQUESTED_DOTNET_SDK_CHANNEL; - const startedAt = Date.now(); - try { - await vscode.commands.executeCommand('dotnet.showAcquisitionLog'); - const acquisition = await vscode.commands.executeCommand<{ dotnetPath?: string } | undefined>( - 'dotnet.acquireGlobalSDK', - { - version: REQUESTED_DOTNET_SDK_CHANNEL, - requestingExtensionId: ext.context.extension.id, - installType: 'global', - }, - ); - telemetryContext.telemetry.measurements.durationMs = Date.now() - startedAt; - const dotnetPath = acquisition?.dotnetPath; - telemetryContext.telemetry.properties.pathReturned = String(!!dotnetPath); - telemetryContext.telemetry.properties.satisfiesMinSdk = dotnetPath - ? String(hasRequiredDotNetSdk(dotnetPath)) - : 'false'; - telemetryContext.telemetry.properties.outcome = dotnetPath ? 'success' : 'noPath'; - return dotnetPath; - } catch (err) { - telemetryContext.telemetry.measurements.durationMs = Date.now() - startedAt; - telemetryContext.telemetry.properties.outcome = 'failure'; - ext.outputChannel.appendLine(`dotnet.acquireGlobalSDK failed: ${String(err)}`); - throw err; - } - }, - ); - return result; -} - -async function promptToInstallDotNetSdk( - context: IActionContext, - node: NoSqlContainerResourceItem | undefined, -): Promise { - const installDotNetSdk = l10n.t('Install .NET SDK'); - const installDotNetTool = l10n.t('Install .NET Install Tool'); - const downloadDotNet = l10n.t('Download .NET SDK'); - const settings = l10n.t('Settings'); - const dotNetInstallToolExtensionId = 'ms-dotnettools.vscode-dotnet-runtime'; - const isDotNetInstallToolInstalled = !!vscode.extensions.getExtension(dotNetInstallToolExtensionId); - const primaryAction = isDotNetInstallToolInstalled ? installDotNetSdk : installDotNetTool; - const selection = await vscode.window.showInformationMessage( - l10n.t( - '.NET SDK {0} or newer is required to install Cosmos DB Shell. Install the .NET SDK, download it manually, or configure an existing Cosmos DB Shell path in settings.', - MIN_DOTNET_SDK_VERSION, - ), - { modal: true }, - primaryAction, - downloadDotNet, - settings, - ); - - const outcome = - selection === installDotNetSdk - ? 'installSdk' - : selection === installDotNetTool - ? 'installTool' - : selection === downloadDotNet - ? 'downloadSdk' - : selection === settings - ? 'settings' - : 'cancelled'; - reportInstallPromptOutcome('installSdk', outcome, { - installToolPresent: String(isDotNetInstallToolInstalled), - }); - - if (selection === installDotNetSdk) { - const dotnetPath = await tryInstallDotNetSdkViaExtension(); - if (dotnetPath && hasRequiredDotNetSdk(dotnetPath)) { - // Chain forward: now that the SDK is available, automatically continue with the - // Cosmos DB Shell install using the freshly-acquired dotnet path so we don't have - // to wait for PATH to be picked up by this VS Code session. - await installAndLaunchCosmosDBShell(context, node, dotnetPath); - } else if (hasRequiredDotNetSdk()) { - await promptToInstallCosmosDBShell(context, node); - } else { - const showOutput = l10n.t('Show Output'); - const failureSelection = await vscode.window.showErrorMessage( - l10n.t( - 'Failed to install .NET SDK {0} or newer. Try downloading it manually from https://dot.net/download.', - MIN_DOTNET_SDK_VERSION, - ), - showOutput, - ); - if (failureSelection === showOutput) { - ext.outputChannel.show(true); - } - } - } else if (selection === installDotNetTool) { - void vscode.commands.executeCommand('workbench.extensions.installExtension', dotNetInstallToolExtensionId); - } else if (selection === downloadDotNet) { - void vscode.env.openExternal(vscode.Uri.parse('https://dot.net/download')); - } else if (selection === settings) { - void vscode.commands.executeCommand('workbench.action.openSettings', 'cosmosDB.shell.path'); - } -} - -async function promptToResolveMissingCosmosDBShell( - context: IActionContext, - node: NoSqlContainerResourceItem | undefined, -): Promise { - if (isCosmosDBShellPathFound()) { - const settings = l10n.t('Settings'); - const selection = await vscode.window.showErrorMessage( - l10n.t( - 'Cosmos DB Shell path is configured but the executable could not be run. Please verify the path in settings.', - ), - settings, - ); - reportInstallPromptOutcome('pathMisconfigured', selection === settings ? 'settings' : 'cancelled'); - if (selection === settings) { - void vscode.commands.executeCommand('workbench.action.openSettings', 'cosmosDB.shell.path'); - } - return; - } - - const sdkOk = hasRequiredDotNetSdk(); - reportInstallPromptOutcome('missingShell', sdkOk ? 'promptInstallShell' : 'promptInstallSdk', { - sdkSatisfiesMin: String(sdkOk), - }); - - if (sdkOk) { - await promptToInstallCosmosDBShell(context, node); - } else { - await promptToInstallDotNetSdk(context, node); - } -} - export async function launchCosmosDBShell(context: IActionContext, node?: NoSqlContainerResourceItem) { const shellInstalled: boolean = isCosmosDBShellInstalled(); // Telemetry: capture launch-shape signals as early as possible so they're attached even // when the install/credential paths bail out before a terminal is created. - const mcpEnabled = SettingsService.getSetting('cosmosDB.shell.MCP.enabled') ?? false; - const mcpPortSetting = SettingsService.getSetting('cosmosDB.shell.MCP.port'); - const mcpPort = (mcpPortSetting ?? 6128).toString(); - const shellPathSetting = SettingsService.getSetting('cosmosDB.shell.path'); + const mcpEnabled = SettingsService.getSetting(SETTING_MCP_ENABLED) ?? false; + const mcpPortSetting = SettingsService.getSetting(SETTING_MCP_PORT); + const mcpPort = (mcpPortSetting ?? DEFAULT_MCP_PORT).toString(); + const shellPathSetting = SettingsService.getSetting(SETTING_SHELL_PATH); context.telemetry.properties.shellInstalled = String(shellInstalled); context.telemetry.properties.shellPathCustom = String(!!shellPathSetting?.trim()); context.telemetry.properties.mcpEnabled = String(mcpEnabled); - context.telemetry.properties.mcpPortDefault = String(mcpPortSetting === undefined || mcpPortSetting === 6128); + context.telemetry.properties.mcpPortDefault = String( + mcpPortSetting === undefined || mcpPortSetting === DEFAULT_MCP_PORT, + ); context.telemetry.properties.authKind = node ? getNodeAuthKind(node) : 'none'; context.telemetry.properties.hasNode = String(!!node); context.telemetry.properties.containerScoped = String(!!node?.model.container); context.telemetry.properties.terminalReused = 'false'; if (!shellInstalled) { - await promptToResolveMissingCosmosDBShell(context, node); + await promptToResolveMissingCosmosDBShell(context, node, launchCosmosDBShell); return; } const command = getCosmosDBShellCommand(); const foundTerminal = vscode.window.terminals.find( - (terminal) => terminal.creationOptions.name === 'Cosmos DB Shell', + (terminal) => terminal.creationOptions.name === COSMOS_DB_SHELL_TERMINAL_NAME, ); + // If another shell terminal is already running, suppress --mcp on the new one: the + // existing process may already own the MCP port, and we don't want to fight over it. const useMcp = mcpEnabled && !foundTerminal; context.telemetry.properties.mcpUsedThisLaunch = String(useMcp); ext.outputChannel.appendLine(`MCP enabled: ${useMcp}, MCP port: ${mcpPort}`); @@ -606,7 +172,7 @@ export async function launchCosmosDBShell(context: IActionContext, node?: NoSqlC } ext.outputChannel.appendLine(`Launching Cosmos DB Shell: ${command} ${args.join(' ')}`); const terminal: vscode.Terminal = vscode.window.createTerminal({ - name: 'Cosmos DB Shell', + name: COSMOS_DB_SHELL_TERMINAL_NAME, shellPath: command, shellArgs: args, }); @@ -673,7 +239,7 @@ export async function launchCosmosDBShell(context: IActionContext, node?: NoSqlC } const terminal: vscode.Terminal = vscode.window.createTerminal({ - name: 'Cosmos DB Shell', + name: COSMOS_DB_SHELL_TERMINAL_NAME, shellPath: command, shellArgs: args, env: Object.keys(env).length > 0 ? env : undefined, @@ -702,13 +268,13 @@ export async function connectCosmosDBShell(context: IActionContext, node?: NoSql context.telemetry.properties.shellVersion = getDetectedCosmosDBShellVersion() ?? 'unknown'; context.telemetry.properties.shellInstalled = String(isCosmosDBShellInstalled()); context.telemetry.properties.shellPathCustom = String( - !!SettingsService.getSetting('cosmosDB.shell.path')?.trim(), + !!SettingsService.getSetting(SETTING_SHELL_PATH)?.trim(), ); - context.telemetry.properties.mcpEnabled = String( - SettingsService.getSetting('cosmosDB.shell.MCP.enabled') ?? false, + context.telemetry.properties.mcpEnabled = String(SettingsService.getSetting(SETTING_MCP_ENABLED) ?? false); + const mcpPortSetting = SettingsService.getSetting(SETTING_MCP_PORT); + context.telemetry.properties.mcpPortDefault = String( + mcpPortSetting === undefined || mcpPortSetting === DEFAULT_MCP_PORT, ); - const mcpPortSetting = SettingsService.getSetting('cosmosDB.shell.MCP.port'); - context.telemetry.properties.mcpPortDefault = String(mcpPortSetting === undefined || mcpPortSetting === 6128); context.telemetry.properties.hasNode = String(!!node); context.telemetry.properties.containerScoped = String(!!node?.model.container); context.telemetry.properties.authKind = node ? getNodeAuthKind(node) : 'none'; @@ -734,8 +300,6 @@ export async function connectCosmosDBShell(context: IActionContext, node?: NoSql if (reusable) { const { terminal } = reusable; context.telemetry.properties.terminalReused = 'true'; - context.telemetry.properties.authKind = getNodeAuthKind(node); - context.telemetry.properties.containerScoped = String(!!node.model.container); terminal.show(); // Always re-issue `connect` before navigating: the shell may have been disconnected // by the user, or previously associated with a different account on a prior reuse. @@ -752,544 +316,3 @@ export async function connectCosmosDBShell(context: IActionContext, node?: NoSql // No reusable terminal (none open, or a different launch-time env is required). await launchCosmosDBShell(context, node); } - -/** - * Classifies the authentication mode required by a node. This determines which env vars - * (if any) the shell process must have been launched with in order to authenticate. - */ -function getNodeAuthKind(node: NoSqlContainerResourceItem): AuthKind { - if (node.model.accountInfo.isEmulator) { - return 'emulator'; - } - if (getCosmosDBShellCredential(node)) { - return 'accountKey'; - } - if (getEntraIdCredential(node)) { - return 'entraId'; - } - if (getManagedIdentityCredential(node)) { - return 'managedIdentity'; - } - return 'none'; -} - -/** Builds a {@link ShellTerminalState} record describing how a shell would be launched for this node. */ -function buildTerminalStateForNode(node: NoSqlContainerResourceItem): ShellTerminalState { - return { - endpoint: node.model.accountInfo.endpoint ?? '', - authKind: getNodeAuthKind(node), - tenantId: getEntraIdCredential(node)?.tenantId, - managedIdentityClientId: getManagedIdentityCredential(node)?.clientId, - }; -} - -/** - * Determines whether an already-running Cosmos DB Shell terminal can host the given node. - * - * Auth modes that need launch-time env vars (account key, Entra ID fallback token) are only - * compatible if the terminal was launched for the *same endpoint* with the *same* auth mode - * (and tenant for Entra ID) — otherwise the baked-in env would be wrong for the new node. - * Auth modes that don't rely on env vars (emulator, managed identity, none) can run in any - * tracked terminal via the interactive `connect` command. - */ -function canReuseTerminalForNode(state: ShellTerminalState, node: NoSqlContainerResourceItem): boolean { - const nodeAuth = getNodeAuthKind(node); - - if (nodeAuth === 'emulator' || nodeAuth === 'managedIdentity' || nodeAuth === 'none') { - return true; - } - - if (state.endpoint !== node.model.accountInfo.endpoint || state.authKind !== nodeAuth) { - return false; - } - - if (nodeAuth === 'entraId') { - const cred = getEntraIdCredential(node); - if (cred?.tenantId !== state.tenantId) { - return false; - } - } - - return true; -} - -/** - * Finds the best tracked Cosmos DB Shell terminal to reuse for the given node, preferring - * terminals already associated with the same endpoint to keep terminal usage stable. - */ -function findReusableTerminalForNode( - node: NoSqlContainerResourceItem, -): { terminal: vscode.Terminal; state: ShellTerminalState } | undefined { - const candidates: Array<{ terminal: vscode.Terminal; state: ShellTerminalState; sameEndpoint: boolean }> = []; - for (const [terminal, state] of terminalStates) { - if (!vscode.window.terminals.includes(terminal)) { - continue; - } - if (!canReuseTerminalForNode(state, node)) { - continue; - } - candidates.push({ - terminal, - state, - sameEndpoint: state.endpoint === node.model.accountInfo.endpoint, - }); - } - candidates.sort((a, b) => Number(b.sameEndpoint) - Number(a.sameEndpoint)); - return candidates[0]; -} - -/** - * Builds the interactive `connect` command that mirrors the CLI `--connect` flag and related - * credential flags, so an already-running Cosmos DB Shell can be attached to a specific account. - */ -function buildInteractiveConnectCommand(node: NoSqlContainerResourceItem, endpoint: string): string { - const parts = ['connect', quoteArg(endpoint)]; - - if (!node.model.accountInfo.isEmulator) { - const entraCredential = getEntraIdCredential(node); - if (entraCredential) { - parts.push('--vscode-credential'); - if (entraCredential.tenantId) { - parts.push('--tenant', quoteArg(entraCredential.tenantId)); - } - } - - const managedIdentityCredential = getManagedIdentityCredential(node); - if (managedIdentityCredential?.clientId) { - parts.push('--managed-identity', quoteArg(managedIdentityCredential.clientId)); - } - } - - return parts.join(' '); -} - -function quoteArg(value: string): string { - return /[\s"']/.test(value) ? `"${value.replace(/"/g, '\\"')}"` : value; -} - -function getCosmosDBShellCredential(node: NoSqlContainerResourceItem): string | undefined { - const credential = node.model.accountInfo.credentials.find((c) => c.type === AuthenticationMethod.accountKey); - return credential?.key; -} - -function getEntraIdCredential(node: NoSqlContainerResourceItem): CosmosDBEntraIdCredential | undefined { - return node.model.accountInfo.credentials.find((c) => c.type === AuthenticationMethod.entraId); -} - -function getManagedIdentityCredential(node: NoSqlContainerResourceItem): CosmosDBManagedIdentityCredential | undefined { - return node.model.accountInfo.credentials.find((c) => c.type === AuthenticationMethod.managedIdentity); -} - -/** - * Obtains an access token from VS Code's authentication session for the Cosmos DB endpoint. - * Used as a fallback token via COSMOSDB_SHELL_TOKEN if VisualStudioCodeCredential fails in the shell. - */ -async function getCosmosDBShellToken( - entraCredential: CosmosDBEntraIdCredential, - endpoint: string, -): Promise { - try { - const endpointUrl = new URL(endpoint); - const scope = `${endpointUrl.origin}${endpointUrl.pathname}.default`; - const token = await getAccessTokenForVSCode(scope, entraCredential.tenantId, { createIfNone: false }); - return token?.token ?? undefined; - } catch { - ext.outputChannel.appendLine('Failed to obtain fallback access token for Cosmos DB Shell'); - return undefined; - } -} - -/** - * Determines if CosmosDBShell is installed. - * - * @returns true, if CosmosDBShell is installed, false otherwise. - */ -export function isCosmosDBShellInstalled(): boolean { - return getCachedShellSupport().installed; -} - -/** - * Returns the version reported by `CosmosDBShell --version` (e.g. `1.2.3` or - * `1.2.3-prerelease.45`), or undefined when the shell is not installed or no - * version could be parsed from its output. - */ -export function getDetectedCosmosDBShellVersion(): string | undefined { - return getCachedShellSupport().version; -} - -/** - * Clears the cached result of {@link isCosmosDBShellInstalled}. - * Call this when the shell path configuration changes or the binary may have been installed/removed. - */ -export function invalidateCosmosDBShellSupportCache(): void { - cosmosDBShellSupportCache.clear(); -} - -type CosmosDBShellSupportInfo = { installed: boolean; version?: string }; - -const cosmosDBShellSupportCache = new Map(); - -function getCachedShellSupport(): CosmosDBShellSupportInfo { - const command = getCosmosDBShellCommand(); - const cached = cosmosDBShellSupportCache.get(command); - if (cached !== undefined) { - return cached; - } - const result = detectCosmosDBShellSupport(command); - cosmosDBShellSupportCache.set(command, result); - return result; -} - -/** - * Extracts a SemVer-like version token (e.g. `1.2.3` or `1.2.3-prerelease.4`) - * from the `--version` output of CosmosDBShell. Returns undefined when no - * recognizable version token is present. - */ -function parseShellVersion(output: string): string | undefined { - const match = output.match(/\b(\d+\.\d+\.\d+(?:[-+][0-9A-Za-z.+-]+)?)\b/); - return match?.[1]; -} - -function detectCosmosDBShellSupport(command: string): CosmosDBShellSupportInfo { - try { - const stdout = child.execFileSync(command, ['--version'], { - windowsHide: true, - env: { - ...process.env, - // CosmosDBShell may use ANSI output libraries (e.g. Spectre.Console). - // When spawned by VS Code, stdio is typically redirected which can confuse - // terminal capability detection. Prefer a plain output mode. - NO_COLOR: '1', - CLICOLOR: '0', - TERM: process.env.TERM ?? 'dumb', - }, - }); - return { installed: true, version: parseShellVersion(stdout.toString('utf8')) }; - } catch (err) { - const anyErr = err as { stdout?: unknown; stderr?: unknown }; - const stdout = - typeof anyErr?.stdout === 'string' - ? anyErr.stdout - : Buffer.isBuffer(anyErr?.stdout) - ? anyErr.stdout.toString('utf8') - : ''; - const stderr = - typeof anyErr?.stderr === 'string' - ? anyErr.stderr - : Buffer.isBuffer(anyErr?.stderr) - ? anyErr.stderr.toString('utf8') - : ''; - - // Workaround: CosmosDBShell may print a valid version string but still exit non-zero - // when ANSI is not available. Treat that as installed. - const combinedOutput = `${stdout}\n${stderr}`; - if (/\bCosmos(?:DB)?Shell\b/i.test(combinedOutput)) { - ext.outputChannel.appendLine( - 'warning: CosmosDBShell "--version" exited non-zero, but returned version output; treating as installed.', - ); - if (stderr.trim().length > 0) { - ext.outputChannel.appendLine(stderr.trim()); - } - return { installed: true, version: parseShellVersion(combinedOutput) }; - } - - ext.outputChannel.appendLine('fail ' + String(err)); - ext.outputChannel.appendLine('while running "' + command + ' --version"'); - if (stdout.trim().length > 0) { - ext.outputChannel.appendLine('stdout: ' + stdout.trim()); - } - if (stderr.trim().length > 0) { - ext.outputChannel.appendLine('stderr: ' + stderr.trim()); - } - return { installed: false }; - } -} -const McpServerName = 'Azure Cosmos DB Shell'; - -function isPortReachable(port: string): Promise { - return new Promise((resolve) => { - const socket = new net.Socket(); - socket.setTimeout(2000); - socket.once('connect', () => { - socket.destroy(); - resolve(true); - }); - socket.once('timeout', () => { - socket.destroy(); - resolve(false); - }); - socket.once('error', () => { - socket.destroy(); - resolve(false); - }); - socket.connect(parseInt(port, 10), CosmosDBShellMcpHost); - }); -} - -function isMcpShellServer(port: string): Promise { - return new Promise((resolve) => { - const req = http.get(`${getCosmosDBShellMcpEndpoint(port)}/sse`, { timeout: 3000 }, (res) => { - const contentType = res.headers['content-type'] ?? ''; - res.destroy(); - resolve(contentType.startsWith('text/event-stream')); - }); - req.once('timeout', () => { - req.destroy(); - resolve(false); - }); - req.once('error', () => { - resolve(false); - }); - }); -} - -function waitForPort( - port: string, - retries: number, - delayMs: number, - token: vscode.CancellationToken, -): Promise { - return new Promise((resolve) => { - let attempt = 0; - const tokenListener = token.onCancellationRequested(() => { - tokenListener.dispose(); - resolve(false); - }); - - const poll = async () => { - if (token.isCancellationRequested) { - tokenListener.dispose(); - resolve(false); - return; - } - if (await isPortReachable(port)) { - tokenListener.dispose(); - resolve(true); - return; - } - attempt++; - if (attempt >= retries) { - tokenListener.dispose(); - resolve(false); - return; - } - setTimeout(() => void poll(), delayMs); - }; - - void poll(); - }); -} - -function showMcpSettingsNotification(message: string, settingKey: string): void { - const settingsLabel = l10n.t('Settings'); - void vscode.window.showWarningMessage(message, settingsLabel).then((selection) => { - if (selection === settingsLabel) { - void vscode.commands.executeCommand('workbench.action.openSettings', settingKey); - } - }); -} - -async function resolveMcpServer( - server: vscode.McpServerDefinition, - mcpPort: string, - token: vscode.CancellationToken, -): Promise { - if (server.label !== McpServerName) { - return server; - } - - const portReachable = await isPortReachable(mcpPort); - - if (portReachable) { - const isShell = await isMcpShellServer(mcpPort); - if (isShell) { - return server; - } - showMcpSettingsNotification( - l10n.t('Port {0} is in use by another process. Configure a different MCP port in settings.', mcpPort), - 'cosmosDB.shell.MCP.port', - ); - throw new Error( - `Port ${mcpPort} is in use by another process that is not the Cosmos DB Shell MCP server. Configure a different port via the "cosmosDB.shell.MCP.port" setting.`, - ); - } - - // No user-facing notifications here: resolve can be invoked by Copilot/VS Code during - // background tool discovery and we don't want to nag users who never asked for Cosmos DB MCP. - // The provider normally hides the server when prerequisites aren't met (see - // provideMcpServerDefinitions); these throws are a safety net for cached definitions. - if (!isCosmosDBShellInstalled()) { - ext.outputChannel.appendLine('MCP resolve: Cosmos DB Shell binary is not installed or not found; skipping.'); - throw new Error( - 'Cosmos DB Shell binary is not installed or not found. The user must install it or configure the "cosmosDB.shell.path" setting.', - ); - } - - const mcpEnabled = SettingsService.getSetting('cosmosDB.shell.MCP.enabled') ?? false; - - if (!mcpEnabled) { - ext.outputChannel.appendLine('MCP resolve: "cosmosDB.shell.MCP.enabled" is disabled; skipping.'); - throw new Error( - 'Cosmos DB Shell MCP is not enabled. The user must enable the "cosmosDB.shell.MCP.enabled" setting and restart the MCP server.', - ); - } - - const existingTerminal = vscode.window.terminals.find((t) => t.creationOptions.name === 'Cosmos DB Shell'); - - if (existingTerminal) { - void vscode.window.showWarningMessage( - l10n.t('The running Cosmos DB Shell was started without MCP. Please close it and try again.'), - ); - throw new Error( - 'A Cosmos DB Shell terminal is already running without MCP support. The user must close it and try again.', - ); - } - - ext.outputChannel.appendLine('MCP resolve: launching Cosmos DB Shell with --mcp'); - await vscode.commands.executeCommand('cosmosDB.launchCosmosDBShell'); - - const ready = await waitForPort(mcpPort, 10, 1000, token); - if (!ready) { - ext.outputChannel.appendLine('MCP resolve: Cosmos DB Shell MCP server did not become reachable in time'); - void vscode.window.showWarningMessage( - l10n.t('Cosmos DB Shell MCP server did not start in time. Check the terminal for errors.'), - ); - throw new Error( - 'Cosmos DB Shell MCP server did not start in time. The user should check the Cosmos DB Shell terminal for errors.', - ); - } - - return server; -} - -export function registerMcpServer(context: vscode.ExtensionContext): void { - try { - const didChangeEmitter = new vscode.EventEmitter(); - - const getMcpPort = (): string => - (SettingsService.getSetting('cosmosDB.shell.MCP.port') ?? 6128).toString(); - - context.subscriptions.push( - vscode.lm.registerMcpServerDefinitionProvider('cosmosDbShellMcpProvider', { - onDidChangeMcpServerDefinitions: didChangeEmitter.event, - provideMcpServerDefinitions: () => { - // Only publish the MCP server when it can actually be used. Otherwise Copilot - // (or any MCP consumer) would call resolveMcpServerDefinition during background - // tool discovery and trigger user-facing prompts even though the user never - // asked for Cosmos DB MCP. The didChangeEmitter below re-fires this when the - // relevant settings or shell path change. - const mcpEnabled = SettingsService.getSetting('cosmosDB.shell.MCP.enabled') ?? false; - if (!mcpEnabled || !isCosmosDBShellInstalled()) { - return []; - } - const mcpPort = getMcpPort(); - return [ - new vscode.McpHttpServerDefinition( - McpServerName, - vscode.Uri.parse(getCosmosDBShellMcpEndpoint(mcpPort)), - { - API_VERSION: '1.0.0', - }, - '1.0.0', - ), - ]; - }, - resolveMcpServerDefinition: (server: vscode.McpServerDefinition, token: vscode.CancellationToken) => { - return resolveMcpServer(server, getMcpPort(), token); - }, - }), - ); - - context.subscriptions.push( - vscode.workspace.onDidChangeConfiguration((event) => { - if (event.affectsConfiguration('cosmosDB.shell.path')) { - invalidateCosmosDBShellSupportCache(); - } - if ( - event.affectsConfiguration('cosmosDB.shell.MCP.port') || - event.affectsConfiguration('cosmosDB.shell.MCP.enabled') || - event.affectsConfiguration('cosmosDB.shell.path') - ) { - didChangeEmitter.fire(); - } - }), - ); - - context.subscriptions.push(didChangeEmitter); - } catch (err) { - ext.outputChannel.appendLine('error while registering MCP server: ' + String(err)); - } -} - -let cosmosDBShellLanguageClient: LanguageClient | undefined; - -export function registerCosmosDBShellLanguageServer(context: vscode.ExtensionContext) { - if (cosmosDBShellLanguageClient || !isCosmosDBShellInstalled()) { - return; - } - - // Path to your LSP server executable - const command = getCosmosDBShellCommand(); - // Adjust argument form depending on the tool’s expectation (--lsp vs -lsp) - const serverArgs = ['--lsp']; - - const serverOptions: ServerOptions = { - run: { - command, - args: serverArgs, - transport: TransportKind.stdio, - }, - debug: { - command, - args: serverArgs, - transport: TransportKind.stdio, - }, - }; - - const clientOptions: LanguageClientOptions = { - documentSelector: [{ scheme: 'file', language: 'cosmosdbshell' }], - synchronize: { - // Watch for related files (adjust pattern as needed) - fileEvents: vscode.workspace.createFileSystemWatcher('**/*.{csh}'), - }, - revealOutputChannelOn: RevealOutputChannelOn.Never, - progressOnInitialization: true, - outputChannelName: l10n.t('Cosmos DB Shell Language Server'), - initializationOptions: { - // Place any feature flags or user settings you want to pass through: - // example: telemetry: true - }, - middleware: { - // Add middleware hooks if needed (e.g. logging, modifications) - }, - }; - - cosmosDBShellLanguageClient = new LanguageClient( - 'cosmosDBShellLanguageServer', - l10n.t('Cosmos DB Shell Language Server'), - serverOptions, - clientOptions, - ); - - context.subscriptions.push({ - dispose: async () => { - if (cosmosDBShellLanguageClient) { - try { - await cosmosDBShellLanguageClient.stop(); - } catch (error) { - console.error('Failed to stop the Cosmos DB Shell language client:', error); - } - cosmosDBShellLanguageClient = undefined; - } - }, - }); - - void cosmosDBShellLanguageClient - .start() - .then(() => { - ext.outputChannel.appendLine('Cosmos DB Shell language server started.'); - }) - .catch((err: unknown) => { - ext.outputChannel.appendLine('Failed to start Cosmos DB Shell language server: ' + String(err)); - }); -} diff --git a/src/cosmosDBShell/constants.ts b/src/cosmosDBShell/constants.ts new file mode 100644 index 000000000..ce4da3212 --- /dev/null +++ b/src/cosmosDBShell/constants.ts @@ -0,0 +1,27 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +/** + * Single source of truth for setting keys, command IDs, and other magic values shared + * across the Cosmos DB Shell modules. Keeping these here ensures renames stay in lockstep + * and that build/lint tooling can flag any stray literal references. + */ + +/** Display name used for every Cosmos DB Shell terminal created by this extension. */ +export const COSMOS_DB_SHELL_TERMINAL_NAME = 'Cosmos DB Shell'; + +/** Default TCP port used by the Cosmos DB Shell MCP server when the user hasn't overridden it. */ +export const DEFAULT_MCP_PORT = 6128; + +/** Label published to the VS Code MCP API for the Cosmos DB Shell MCP server. */ +export const MCP_SERVER_NAME = 'Azure Cosmos DB Shell'; + +// --- VS Code setting keys (must match the contributions in package.json) --- +export const SETTING_SHELL_PATH = 'cosmosDB.shell.path'; +export const SETTING_MCP_ENABLED = 'cosmosDB.shell.MCP.enabled'; +export const SETTING_MCP_PORT = 'cosmosDB.shell.MCP.port'; + +// --- VS Code command IDs --- +export const COMMAND_LAUNCH_COSMOS_DB_SHELL = 'cosmosDB.launchCosmosDBShell'; diff --git a/src/cosmosDBShell/install/dotNetSdk.test.ts b/src/cosmosDBShell/install/dotNetSdk.test.ts new file mode 100644 index 000000000..884ad1f5f --- /dev/null +++ b/src/cosmosDBShell/install/dotNetSdk.test.ts @@ -0,0 +1,138 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import * as child from 'child_process'; +import { beforeEach, describe, expect, it, vi, type Mock } from 'vitest'; +import { + MIN_DOTNET_SDK_VERSION, + REQUESTED_DOTNET_SDK_CHANNEL, + compareDotNetVersions, + getInstalledDotNetSdkVersions, + hasRequiredDotNetSdk, +} from './dotNetSdk'; + +vi.mock('child_process', () => ({ + execFileSync: vi.fn(), +})); + +vi.mock('../../extensionVariables', () => ({ + ext: { + outputChannel: { + appendLine: vi.fn(), + }, + }, +})); + +// @microsoft/vscode-azext-utils transitively pulls in vscode via CJS requires. Stub the only +// symbol dotNetSdk uses (callWithTelemetryAndErrorHandling) so the module loads cleanly. +vi.mock('@microsoft/vscode-azext-utils', () => ({ + callWithTelemetryAndErrorHandling: vi.fn(), +})); + +describe('dotNetSdk.compareDotNetVersions', () => { + it('compares numeric major / minor / patch components', () => { + expect(compareDotNetVersions('10.0.0', '9.0.0')).toBeGreaterThan(0); + expect(compareDotNetVersions('9.0.0', '10.0.0')).toBeLessThan(0); + expect(compareDotNetVersions('10.0.203', '10.0.100')).toBeGreaterThan(0); + expect(compareDotNetVersions('10.1.0', '10.0.999')).toBeGreaterThan(0); + }); + + it('returns 0 for equal versions', () => { + expect(compareDotNetVersions('10.0.203', '10.0.203')).toBe(0); + }); + + it('ignores pre-release / build suffixes after "-"', () => { + expect(compareDotNetVersions('9.0.100-rc.1', '9.0.100')).toBe(0); + expect(compareDotNetVersions('10.0.203-preview.1', '10.0.203-rc.2')).toBe(0); + }); + + it('treats missing components as zero', () => { + expect(compareDotNetVersions('10', '10.0.0')).toBe(0); + expect(compareDotNetVersions('10.0.0', '10')).toBe(0); + expect(compareDotNetVersions('10.1', '10.0.999')).toBeGreaterThan(0); + }); + + it('treats non-numeric components as zero', () => { + expect(compareDotNetVersions('abc.def.ghi', '0.0.0')).toBe(0); + }); +}); + +describe('dotNetSdk.REQUESTED_DOTNET_SDK_CHANNEL', () => { + it('is the "major.minor" prefix of MIN_DOTNET_SDK_VERSION', () => { + const expected = MIN_DOTNET_SDK_VERSION.split('.').slice(0, 2).join('.'); + expect(REQUESTED_DOTNET_SDK_CHANNEL).toBe(expected); + }); +}); + +describe('dotNetSdk.getInstalledDotNetSdkVersions', () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it('parses one version per non-empty line of "dotnet --list-sdks" output', () => { + (child.execFileSync as Mock).mockReturnValue( + '8.0.404 [C:\\Program Files\\dotnet\\sdk]\r\n10.0.203 [C:\\Program Files\\dotnet\\sdk]\r\n', + ); + expect(getInstalledDotNetSdkVersions()).toEqual(['8.0.404', '10.0.203']); + }); + + it('accepts Buffer return values from execFileSync', () => { + (child.execFileSync as Mock).mockReturnValue(Buffer.from('9.0.100 [x]\n')); + expect(getInstalledDotNetSdkVersions()).toEqual(['9.0.100']); + }); + + it('returns an empty array when execFileSync throws (dotnet not on PATH)', () => { + (child.execFileSync as Mock).mockImplementation(() => { + throw new Error('not found'); + }); + expect(getInstalledDotNetSdkVersions()).toEqual([]); + }); + + it('skips blank lines', () => { + (child.execFileSync as Mock).mockReturnValue('\n9.0.100 [x]\n\n'); + expect(getInstalledDotNetSdkVersions()).toEqual(['9.0.100']); + }); + + it('invokes the supplied dotnetPath instead of "dotnet" when provided', () => { + (child.execFileSync as Mock).mockReturnValue('10.0.203 [x]\n'); + getInstalledDotNetSdkVersions('/custom/dotnet'); + expect((child.execFileSync as Mock).mock.calls[0][0]).toBe('/custom/dotnet'); + }); + + it('defaults to "dotnet" when no dotnetPath is provided', () => { + (child.execFileSync as Mock).mockReturnValue('10.0.203 [x]\n'); + getInstalledDotNetSdkVersions(); + expect((child.execFileSync as Mock).mock.calls[0][0]).toBe('dotnet'); + }); +}); + +describe('dotNetSdk.hasRequiredDotNetSdk', () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it('returns true when at least one installed SDK >= MIN_DOTNET_SDK_VERSION', () => { + (child.execFileSync as Mock).mockReturnValue(`8.0.404 [x]\n${MIN_DOTNET_SDK_VERSION} [x]\n`); + expect(hasRequiredDotNetSdk()).toBe(true); + }); + + it('returns false when all installed SDKs are below MIN_DOTNET_SDK_VERSION', () => { + (child.execFileSync as Mock).mockReturnValue('8.0.404 [x]\n9.0.100 [x]\n'); + expect(hasRequiredDotNetSdk()).toBe(false); + }); + + it('returns false when no SDKs are installed (execFileSync throws)', () => { + (child.execFileSync as Mock).mockImplementation(() => { + throw new Error('not found'); + }); + expect(hasRequiredDotNetSdk()).toBe(false); + }); + + it('forwards the supplied dotnetPath to getInstalledDotNetSdkVersions', () => { + (child.execFileSync as Mock).mockReturnValue(`${MIN_DOTNET_SDK_VERSION} [x]\n`); + hasRequiredDotNetSdk('/custom/dotnet'); + expect((child.execFileSync as Mock).mock.calls[0][0]).toBe('/custom/dotnet'); + }); +}); diff --git a/src/cosmosDBShell/install/dotNetSdk.ts b/src/cosmosDBShell/install/dotNetSdk.ts new file mode 100644 index 000000000..2d73d0f96 --- /dev/null +++ b/src/cosmosDBShell/install/dotNetSdk.ts @@ -0,0 +1,129 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +/** + * .NET SDK detection and acquisition helpers used by the Cosmos DB Shell install flow. + * + * The CosmosDBShell tool is distributed as a `dotnet tool` global package, so installing + * it (and keeping the user happy when it's missing) requires both a `dotnet` CLI on PATH + * *and* an SDK whose version satisfies {@link MIN_DOTNET_SDK_VERSION}. + */ +import { callWithTelemetryAndErrorHandling, type IActionContext } from '@microsoft/vscode-azext-utils'; +import * as child from 'child_process'; +import * as vscode from 'vscode'; +import { ext } from '../../extensionVariables'; + +/** + * Minimum .NET SDK version required to install the Cosmos DB Shell global tool. + * Bumping this constant also updates the version requested via + * `dotnet.acquireGlobalSDK` and the user-facing prompt copy. + */ +export const MIN_DOTNET_SDK_VERSION = '10.0.203'; + +/** + * Channel (`major.minor`) requested from `dotnet.acquireGlobalSDK`. Using a + * channel rather than {@link MIN_DOTNET_SDK_VERSION} lets the .NET Install Tool + * resolve to the latest available patch on that channel while still satisfying + * the floor enforced by {@link hasRequiredDotNetSdk}. + */ +export const REQUESTED_DOTNET_SDK_CHANNEL = MIN_DOTNET_SDK_VERSION.split('.').slice(0, 2).join('.'); + +/** + * Compares two .NET SDK version strings (e.g. `10.0.203`, `9.0.100-rc.1`) by + * numeric major / minor / patch components. Any pre-release / build metadata + * suffix (after `-`) is ignored. Returns a negative number when `a < b`, zero + * when equal, and a positive number when `a > b`. + */ +export function compareDotNetVersions(a: string, b: string): number { + const parse = (v: string): number[] => + v + .split('-')[0] + .split('.') + .map((part) => Number.parseInt(part, 10) || 0); + const av = parse(a); + const bv = parse(b); + const len = Math.max(av.length, bv.length); + for (let i = 0; i < len; i++) { + const diff = (av[i] ?? 0) - (bv[i] ?? 0); + if (diff !== 0) { + return diff; + } + } + return 0; +} + +/** + * Runs `dotnet --list-sdks` and returns the parsed version strings. The CLI + * output format is ` []` per line. Returns an empty + * array when `dotnet` is not on PATH or the call fails. + */ +export function getInstalledDotNetSdkVersions(dotnetPath?: string): string[] { + try { + const output = child.execFileSync(dotnetPath ?? 'dotnet', ['--list-sdks'], { + windowsHide: true, + stdio: 'pipe', + }); + return output + .toString('utf8') + .split(/\r?\n/) + .map((line) => line.trim()) + .filter((line) => line.length > 0) + .map((line) => line.split(/\s+/, 1)[0]); + } catch { + return []; + } +} + +export function hasRequiredDotNetSdk(dotnetPath?: string): boolean { + return getInstalledDotNetSdkVersions(dotnetPath).some( + (version) => compareDotNetVersions(version, MIN_DOTNET_SDK_VERSION) >= 0, + ); +} + +/** + * Awaits the `.NET Install Tool` SDK acquisition command, requesting a global + * install of the latest patch available on {@link REQUESTED_DOTNET_SDK_CHANNEL}. + * Returns the resolved `dotnet` executable path on success, or undefined when + * the acquisition failed or did not return a path. + * + * Uses `dotnet.acquireGlobalSDK` rather than `dotnet.acquireGlobalSDKPublic` + * because the public variant ignores the supplied `version` and prompts the + * user with its own recommended version instead. + */ +export async function tryInstallDotNetSdkViaExtension(): Promise { + const result = await callWithTelemetryAndErrorHandling( + 'cosmosDB.cosmosDBShell.install.dotnetSdk', + async (telemetryContext: IActionContext) => { + telemetryContext.errorHandling.suppressDisplay = true; + telemetryContext.telemetry.properties.requestedChannel = REQUESTED_DOTNET_SDK_CHANNEL; + const startedAt = Date.now(); + try { + await vscode.commands.executeCommand('dotnet.showAcquisitionLog'); + const acquisition = await vscode.commands.executeCommand<{ dotnetPath?: string } | undefined>( + 'dotnet.acquireGlobalSDK', + { + version: REQUESTED_DOTNET_SDK_CHANNEL, + requestingExtensionId: ext.context.extension.id, + installType: 'global', + }, + ); + telemetryContext.telemetry.measurements.durationMs = Date.now() - startedAt; + const dotnetPath = acquisition?.dotnetPath; + telemetryContext.telemetry.properties.pathReturned = String(!!dotnetPath); + telemetryContext.telemetry.properties.satisfiesMinSdk = dotnetPath + ? String(hasRequiredDotNetSdk(dotnetPath)) + : 'false'; + telemetryContext.telemetry.properties.outcome = dotnetPath ? 'success' : 'noPath'; + return dotnetPath; + } catch (err) { + telemetryContext.telemetry.measurements.durationMs = Date.now() - startedAt; + telemetryContext.telemetry.properties.outcome = 'failure'; + ext.outputChannel.appendLine(`dotnet.acquireGlobalSDK failed: ${String(err)}`); + throw err; + } + }, + ); + return result; +} diff --git a/src/cosmosDBShell/install/installPrompts.ts b/src/cosmosDBShell/install/installPrompts.ts new file mode 100644 index 000000000..c44867149 --- /dev/null +++ b/src/cosmosDBShell/install/installPrompts.ts @@ -0,0 +1,313 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +/** + * User-facing install/repair flow for the CosmosDBShell `dotnet tool` package + * and its .NET SDK prerequisite. + * + * All prompts emit a `cosmosDB.cosmosDBShell.install.prompt` telemetry event so the + * install funnel can be measured without depending on (localized) button labels. + * + * A `launchShell` callback is injected to avoid a circular dependency with the + * main extension module that owns the actual launch flow. + */ +import { callWithTelemetryAndErrorHandling, type IActionContext } from '@microsoft/vscode-azext-utils'; +import * as l10n from '@vscode/l10n'; +import * as child from 'child_process'; +import * as vscode from 'vscode'; +import { ext } from '../../extensionVariables'; +import { type NoSqlContainerResourceItem } from '../../tree/nosql/NoSqlContainerResourceItem'; +import { SETTING_SHELL_PATH } from '../constants'; +import { isCosmosDBShellPathFound } from '../shellCommand'; +import { isCosmosDBShellInstalled } from '../shellSupportCache'; +import { MIN_DOTNET_SDK_VERSION, hasRequiredDotNetSdk, tryInstallDotNetSdkViaExtension } from './dotNetSdk'; + +/** Callback signature used by the install flow to resume the original launch action after install. */ +export type LaunchShellFn = (context: IActionContext, node: NoSqlContainerResourceItem | undefined) => Promise; + +/** + * Runs `dotnet tool install --global CosmosDBShell --prerelease` with a progress + * notification, streaming output to the extension output channel. Returns true + * when the process exits with code 0. + */ +async function installCosmosDBShellWithDotNetTool(dotnetPath?: string): Promise { + const result = await callWithTelemetryAndErrorHandling( + 'cosmosDB.cosmosDBShell.install.tool', + async (telemetryContext: IActionContext) => { + telemetryContext.errorHandling.suppressDisplay = true; + telemetryContext.telemetry.properties.dotnetPathProvided = String(!!dotnetPath); + const startedAt = Date.now(); + const outcome = await vscode.window.withProgress( + { + location: vscode.ProgressLocation.Notification, + title: l10n.t('Installing Cosmos DB Shell…'), + cancellable: true, + }, + async (_progress, token) => { + ext.outputChannel.show(true); + const dotnetExe = dotnetPath ?? 'dotnet'; + ext.outputChannel.appendLine(`> ${dotnetExe} tool install --global CosmosDBShell --prerelease`); + + return new Promise<{ success: boolean; exitCode: number | null; cancelled: boolean }>((resolve) => { + let cancelled = false; + const proc = child.spawn( + dotnetExe, + ['tool', 'install', '--global', 'CosmosDBShell', '--prerelease'], + { windowsHide: true, shell: false }, + ); + + token.onCancellationRequested(() => { + cancelled = true; + proc.kill(); + }); + + proc.stdout?.on('data', (data: Buffer) => { + ext.outputChannel.append(data.toString('utf8')); + }); + proc.stderr?.on('data', (data: Buffer) => { + ext.outputChannel.append(data.toString('utf8')); + }); + proc.on('error', (err) => { + ext.outputChannel.appendLine(`Failed to start dotnet: ${err.message}`); + resolve({ success: false, exitCode: null, cancelled }); + }); + proc.on('close', (code) => { + ext.outputChannel.appendLine(`\nProcess exited with code ${code}.`); + resolve({ success: code === 0, exitCode: code, cancelled }); + }); + }); + }, + ); + telemetryContext.telemetry.measurements.durationMs = Date.now() - startedAt; + telemetryContext.telemetry.properties.exitCode = + outcome.exitCode === null ? 'null' : String(outcome.exitCode); + telemetryContext.telemetry.properties.cancelled = String(outcome.cancelled); + telemetryContext.telemetry.properties.outcome = outcome.cancelled + ? 'cancelled' + : outcome.success + ? 'success' + : 'failure'; + return outcome.success; + }, + ); + return result ?? false; +} + +/** + * Fires a `cosmosDB.cosmosDBShell.install.prompt` telemetry event with the + * given prompt identifier and user selection. Used to measure the install + * funnel without depending on localized button labels. + */ +function reportInstallPromptOutcome( + promptKind: + | 'missingShell' + | 'installShell' + | 'installSdk' + | 'pathMisconfigured' + | 'reloadAfterInstall' + | 'installFailure', + selection: string, + extraProperties?: Record, +): void { + void callWithTelemetryAndErrorHandling( + 'cosmosDB.cosmosDBShell.install.prompt', + (telemetryContext: IActionContext) => { + telemetryContext.errorHandling.suppressDisplay = true; + telemetryContext.telemetry.properties.promptKind = promptKind; + telemetryContext.telemetry.properties.selection = selection; + if (extraProperties) { + for (const [k, v] of Object.entries(extraProperties)) { + telemetryContext.telemetry.properties[k] = v; + } + } + }, + ); +} + +/** + * Prompts the user to install Cosmos DB Shell via `dotnet tool install`, and on + * success automatically continues the original launch flow. + */ +async function promptToInstallCosmosDBShell( + context: IActionContext, + node: NoSqlContainerResourceItem | undefined, + launchShell: LaunchShellFn, +): Promise { + const install = l10n.t('Install'); + const settings = l10n.t('Settings'); + const selection = await vscode.window.showInformationMessage( + l10n.t( + 'Cosmos DB Shell is not installed. Install it now using `dotnet tool install --global CosmosDBShell --prerelease`?', + ), + { modal: true }, + install, + settings, + ); + + const outcome = selection === install ? 'install' : selection === settings ? 'settings' : 'cancelled'; + reportInstallPromptOutcome('installShell', outcome); + + if (selection === settings) { + void vscode.commands.executeCommand('workbench.action.openSettings', SETTING_SHELL_PATH); + return; + } + if (selection !== install) { + return; + } + + await installAndLaunchCosmosDBShell(context, node, launchShell); +} + +/** + * Runs the `dotnet tool install` for Cosmos DB Shell, then either reloads the window + * (if PATH hasn't picked up the new tool yet) or auto-launches the shell to continue + * the user's original action. Used both by the explicit install prompt and by the + * auto-chain after a fresh .NET SDK acquisition. + */ +async function installAndLaunchCosmosDBShell( + context: IActionContext, + node: NoSqlContainerResourceItem | undefined, + launchShell: LaunchShellFn, + dotnetPath?: string, +): Promise { + const success = await installCosmosDBShellWithDotNetTool(dotnetPath); + if (!success) { + const showOutput = l10n.t('Show Output'); + const failureSelection = await vscode.window.showErrorMessage( + l10n.t('Failed to install Cosmos DB Shell. See the output for details.'), + showOutput, + ); + reportInstallPromptOutcome('installFailure', failureSelection === showOutput ? 'showOutput' : 'dismissed'); + if (failureSelection === showOutput) { + ext.outputChannel.show(true); + } + return; + } + + // On a brand-new install the user's PATH may not yet include `~/.dotnet/tools` + // in the current VS Code session. If we still can't resolve the shell, ask to reload. + if (!isCosmosDBShellInstalled()) { + const reload = l10n.t('Reload Window'); + const reloadSelection = await vscode.window.showInformationMessage( + l10n.t( + 'Cosmos DB Shell was installed, but its location is not yet on PATH for this VS Code window. Reload the window to pick it up.', + ), + reload, + ); + reportInstallPromptOutcome('reloadAfterInstall', reloadSelection === reload ? 'reload' : 'cancelled'); + if (reloadSelection === reload) { + void vscode.commands.executeCommand('workbench.action.reloadWindow'); + } + return; + } + + // Auto-relaunch with the original node so the user lands where they intended. + await launchShell(context, node); +} + +async function promptToInstallDotNetSdk( + context: IActionContext, + node: NoSqlContainerResourceItem | undefined, + launchShell: LaunchShellFn, +): Promise { + const installDotNetSdk = l10n.t('Install .NET SDK'); + const installDotNetTool = l10n.t('Install .NET Install Tool'); + const downloadDotNet = l10n.t('Download .NET SDK'); + const settings = l10n.t('Settings'); + const dotNetInstallToolExtensionId = 'ms-dotnettools.vscode-dotnet-runtime'; + const isDotNetInstallToolInstalled = !!vscode.extensions.getExtension(dotNetInstallToolExtensionId); + const primaryAction = isDotNetInstallToolInstalled ? installDotNetSdk : installDotNetTool; + const selection = await vscode.window.showInformationMessage( + l10n.t( + '.NET SDK {0} or newer is required to install Cosmos DB Shell. Install the .NET SDK, download it manually, or configure an existing Cosmos DB Shell path in settings.', + MIN_DOTNET_SDK_VERSION, + ), + { modal: true }, + primaryAction, + downloadDotNet, + settings, + ); + + const outcome = + selection === installDotNetSdk + ? 'installSdk' + : selection === installDotNetTool + ? 'installTool' + : selection === downloadDotNet + ? 'downloadSdk' + : selection === settings + ? 'settings' + : 'cancelled'; + reportInstallPromptOutcome('installSdk', outcome, { + installToolPresent: String(isDotNetInstallToolInstalled), + }); + + if (selection === installDotNetSdk) { + const dotnetPath = await tryInstallDotNetSdkViaExtension(); + if (dotnetPath && hasRequiredDotNetSdk(dotnetPath)) { + // Chain forward: now that the SDK is available, automatically continue with the + // Cosmos DB Shell install using the freshly-acquired dotnet path so we don't have + // to wait for PATH to be picked up by this VS Code session. + await installAndLaunchCosmosDBShell(context, node, launchShell, dotnetPath); + } else if (hasRequiredDotNetSdk()) { + await promptToInstallCosmosDBShell(context, node, launchShell); + } else { + const showOutput = l10n.t('Show Output'); + const failureSelection = await vscode.window.showErrorMessage( + l10n.t( + 'Failed to install .NET SDK {0} or newer. Try downloading it manually from https://dot.net/download.', + MIN_DOTNET_SDK_VERSION, + ), + showOutput, + ); + if (failureSelection === showOutput) { + ext.outputChannel.show(true); + } + } + } else if (selection === installDotNetTool) { + void vscode.commands.executeCommand('workbench.extensions.installExtension', dotNetInstallToolExtensionId); + } else if (selection === downloadDotNet) { + void vscode.env.openExternal(vscode.Uri.parse('https://dot.net/download')); + } else if (selection === settings) { + void vscode.commands.executeCommand('workbench.action.openSettings', SETTING_SHELL_PATH); + } +} + +/** + * Top-level entry point: when the launch flow discovers the shell isn't installed, + * branches to the appropriate prompt depending on whether the configured path is + * broken, the .NET SDK is missing, or just the tool itself is missing. + */ +export async function promptToResolveMissingCosmosDBShell( + context: IActionContext, + node: NoSqlContainerResourceItem | undefined, + launchShell: LaunchShellFn, +): Promise { + if (isCosmosDBShellPathFound()) { + const settings = l10n.t('Settings'); + const selection = await vscode.window.showErrorMessage( + l10n.t( + 'Cosmos DB Shell path is configured but the executable could not be run. Please verify the path in settings.', + ), + settings, + ); + reportInstallPromptOutcome('pathMisconfigured', selection === settings ? 'settings' : 'cancelled'); + if (selection === settings) { + void vscode.commands.executeCommand('workbench.action.openSettings', SETTING_SHELL_PATH); + } + return; + } + + const sdkOk = hasRequiredDotNetSdk(); + reportInstallPromptOutcome('missingShell', sdkOk ? 'promptInstallShell' : 'promptInstallSdk', { + sdkSatisfiesMin: String(sdkOk), + }); + + if (sdkOk) { + await promptToInstallCosmosDBShell(context, node, launchShell); + } else { + await promptToInstallDotNetSdk(context, node, launchShell); + } +} diff --git a/src/cosmosDBShell/languageServer.ts b/src/cosmosDBShell/languageServer.ts new file mode 100644 index 000000000..c367371f1 --- /dev/null +++ b/src/cosmosDBShell/languageServer.ts @@ -0,0 +1,94 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +/** + * Starts the Cosmos DB Shell language server (LSP over stdio) using the same CosmosDBShell + * binary that powers the interactive terminal, invoked with `--lsp`. + */ +import * as l10n from '@vscode/l10n'; +import * as vscode from 'vscode'; +import { + LanguageClient, + RevealOutputChannelOn, + TransportKind, + type LanguageClientOptions, + type ServerOptions, +} from 'vscode-languageclient/node'; +import { ext } from '../extensionVariables'; +import { getCosmosDBShellCommand } from './shellCommand'; +import { isCosmosDBShellInstalled } from './shellSupportCache'; + +let cosmosDBShellLanguageClient: LanguageClient | undefined; + +export function registerCosmosDBShellLanguageServer(context: vscode.ExtensionContext) { + if (cosmosDBShellLanguageClient || !isCosmosDBShellInstalled()) { + return; + } + + // Path to the LSP server executable + const command = getCosmosDBShellCommand(); + // Adjust argument form depending on the tool's expectation (--lsp vs -lsp) + const serverArgs = ['--lsp']; + + const serverOptions: ServerOptions = { + run: { + command, + args: serverArgs, + transport: TransportKind.stdio, + }, + debug: { + command, + args: serverArgs, + transport: TransportKind.stdio, + }, + }; + + const clientOptions: LanguageClientOptions = { + documentSelector: [{ scheme: 'file', language: 'cosmosdbshell' }], + synchronize: { + // Watch for related files (adjust pattern as needed) + fileEvents: vscode.workspace.createFileSystemWatcher('**/*.{csh}'), + }, + revealOutputChannelOn: RevealOutputChannelOn.Never, + progressOnInitialization: true, + outputChannelName: l10n.t('Cosmos DB Shell Language Server'), + initializationOptions: { + // Place any feature flags or user settings you want to pass through: + // example: telemetry: true + }, + middleware: { + // Add middleware hooks if needed (e.g. logging, modifications) + }, + }; + + cosmosDBShellLanguageClient = new LanguageClient( + 'cosmosDBShellLanguageServer', + l10n.t('Cosmos DB Shell Language Server'), + serverOptions, + clientOptions, + ); + + context.subscriptions.push({ + dispose: async () => { + if (cosmosDBShellLanguageClient) { + try { + await cosmosDBShellLanguageClient.stop(); + } catch (error) { + console.error('Failed to stop the Cosmos DB Shell language client:', error); + } + cosmosDBShellLanguageClient = undefined; + } + }, + }); + + void cosmosDBShellLanguageClient + .start() + .then(() => { + ext.outputChannel.appendLine('Cosmos DB Shell language server started.'); + }) + .catch((err: unknown) => { + ext.outputChannel.appendLine('Failed to start Cosmos DB Shell language server: ' + String(err)); + }); +} diff --git a/src/cosmosDBShell/mcpProvider.ts b/src/cosmosDBShell/mcpProvider.ts new file mode 100644 index 000000000..db7324097 --- /dev/null +++ b/src/cosmosDBShell/mcpProvider.ts @@ -0,0 +1,245 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +/** + * VS Code MCP server provider for the Cosmos DB Shell. Publishes a single + * {@link vscode.McpHttpServerDefinition} when the shell binary is installed and the + * `cosmosDB.shell.MCP.enabled` setting is on, and lazily launches the shell with `--mcp` + * when VS Code or Copilot asks to resolve the server. + */ +import * as l10n from '@vscode/l10n'; +import * as http from 'http'; +import * as net from 'net'; +import * as vscode from 'vscode'; +import { ext } from '../extensionVariables'; +import { SettingsService } from '../services/SettingsService'; +import { + COMMAND_LAUNCH_COSMOS_DB_SHELL, + COSMOS_DB_SHELL_TERMINAL_NAME, + DEFAULT_MCP_PORT, + MCP_SERVER_NAME, + SETTING_MCP_ENABLED, + SETTING_MCP_PORT, + SETTING_SHELL_PATH, +} from './constants'; +import { CosmosDBShellMcpHost, getCosmosDBShellMcpEndpoint } from './cosmosDBShellMcpEndpoint'; +import { invalidateCosmosDBShellSupportCache, isCosmosDBShellInstalled } from './shellSupportCache'; + +function isPortReachable(port: string): Promise { + return new Promise((resolve) => { + const socket = new net.Socket(); + socket.setTimeout(2000); + socket.once('connect', () => { + socket.destroy(); + resolve(true); + }); + socket.once('timeout', () => { + socket.destroy(); + resolve(false); + }); + socket.once('error', () => { + socket.destroy(); + resolve(false); + }); + socket.connect(parseInt(port, 10), CosmosDBShellMcpHost); + }); +} + +function isMcpShellServer(port: string): Promise { + return new Promise((resolve) => { + const req = http.get(`${getCosmosDBShellMcpEndpoint(port)}/sse`, { timeout: 3000 }, (res) => { + const contentType = res.headers['content-type'] ?? ''; + res.destroy(); + resolve(contentType.startsWith('text/event-stream')); + }); + req.once('timeout', () => { + req.destroy(); + resolve(false); + }); + req.once('error', () => { + resolve(false); + }); + }); +} + +function waitForPort( + port: string, + retries: number, + delayMs: number, + token: vscode.CancellationToken, +): Promise { + return new Promise((resolve) => { + let attempt = 0; + const tokenListener = token.onCancellationRequested(() => { + tokenListener.dispose(); + resolve(false); + }); + + const poll = async () => { + if (token.isCancellationRequested) { + tokenListener.dispose(); + resolve(false); + return; + } + if (await isPortReachable(port)) { + tokenListener.dispose(); + resolve(true); + return; + } + attempt++; + if (attempt >= retries) { + tokenListener.dispose(); + resolve(false); + return; + } + setTimeout(() => void poll(), delayMs); + }; + + void poll(); + }); +} + +function showMcpSettingsNotification(message: string, settingKey: string): void { + const settingsLabel = l10n.t('Settings'); + void vscode.window.showWarningMessage(message, settingsLabel).then((selection) => { + if (selection === settingsLabel) { + void vscode.commands.executeCommand('workbench.action.openSettings', settingKey); + } + }); +} + +async function resolveMcpServer( + server: vscode.McpServerDefinition, + mcpPort: string, + token: vscode.CancellationToken, +): Promise { + if (server.label !== MCP_SERVER_NAME) { + return server; + } + + const portReachable = await isPortReachable(mcpPort); + + if (portReachable) { + const isShell = await isMcpShellServer(mcpPort); + if (isShell) { + return server; + } + showMcpSettingsNotification( + l10n.t('Port {0} is in use by another process. Configure a different MCP port in settings.', mcpPort), + SETTING_MCP_PORT, + ); + throw new Error( + `Port ${mcpPort} is in use by another process that is not the Cosmos DB Shell MCP server. Configure a different port via the "${SETTING_MCP_PORT}" setting.`, + ); + } + + // No user-facing notifications here: resolve can be invoked by Copilot/VS Code during + // background tool discovery and we don't want to nag users who never asked for Cosmos DB MCP. + // The provider normally hides the server when prerequisites aren't met (see + // provideMcpServerDefinitions); these throws are a safety net for cached definitions. + if (!isCosmosDBShellInstalled()) { + ext.outputChannel.appendLine('MCP resolve: Cosmos DB Shell binary is not installed or not found; skipping.'); + throw new Error( + `Cosmos DB Shell binary is not installed or not found. The user must install it or configure the "${SETTING_SHELL_PATH}" setting.`, + ); + } + + const mcpEnabled = SettingsService.getSetting(SETTING_MCP_ENABLED) ?? false; + + if (!mcpEnabled) { + ext.outputChannel.appendLine(`MCP resolve: "${SETTING_MCP_ENABLED}" is disabled; skipping.`); + throw new Error( + `Cosmos DB Shell MCP is not enabled. The user must enable the "${SETTING_MCP_ENABLED}" setting and restart the MCP server.`, + ); + } + + const existingTerminal = vscode.window.terminals.find( + (t) => t.creationOptions.name === COSMOS_DB_SHELL_TERMINAL_NAME, + ); + + if (existingTerminal) { + void vscode.window.showWarningMessage( + l10n.t('The running Cosmos DB Shell was started without MCP. Please close it and try again.'), + ); + throw new Error( + 'A Cosmos DB Shell terminal is already running without MCP support. The user must close it and try again.', + ); + } + + ext.outputChannel.appendLine('MCP resolve: launching Cosmos DB Shell with --mcp'); + await vscode.commands.executeCommand(COMMAND_LAUNCH_COSMOS_DB_SHELL); + + const ready = await waitForPort(mcpPort, 10, 1000, token); + if (!ready) { + ext.outputChannel.appendLine('MCP resolve: Cosmos DB Shell MCP server did not become reachable in time'); + void vscode.window.showWarningMessage( + l10n.t('Cosmos DB Shell MCP server did not start in time. Check the terminal for errors.'), + ); + throw new Error( + 'Cosmos DB Shell MCP server did not start in time. The user should check the Cosmos DB Shell terminal for errors.', + ); + } + + return server; +} + +export function registerMcpServer(context: vscode.ExtensionContext): void { + try { + const didChangeEmitter = new vscode.EventEmitter(); + + const getMcpPort = (): string => + (SettingsService.getSetting(SETTING_MCP_PORT) ?? DEFAULT_MCP_PORT).toString(); + + context.subscriptions.push( + vscode.lm.registerMcpServerDefinitionProvider('cosmosDbShellMcpProvider', { + onDidChangeMcpServerDefinitions: didChangeEmitter.event, + provideMcpServerDefinitions: () => { + // Only publish the MCP server when it can actually be used. Otherwise Copilot + // (or any MCP consumer) would call resolveMcpServerDefinition during background + // tool discovery and trigger user-facing prompts even though the user never + // asked for Cosmos DB MCP. The didChangeEmitter below re-fires this when the + // relevant settings or shell path change. + const mcpEnabled = SettingsService.getSetting(SETTING_MCP_ENABLED) ?? false; + if (!mcpEnabled || !isCosmosDBShellInstalled()) { + return []; + } + const mcpPort = getMcpPort(); + return [ + new vscode.McpHttpServerDefinition( + MCP_SERVER_NAME, + vscode.Uri.parse(getCosmosDBShellMcpEndpoint(mcpPort)), + { + API_VERSION: '1.0.0', + }, + '1.0.0', + ), + ]; + }, + resolveMcpServerDefinition: (server: vscode.McpServerDefinition, token: vscode.CancellationToken) => { + return resolveMcpServer(server, getMcpPort(), token); + }, + }), + ); + + context.subscriptions.push( + vscode.workspace.onDidChangeConfiguration((event) => { + if (event.affectsConfiguration(SETTING_SHELL_PATH)) { + invalidateCosmosDBShellSupportCache(); + } + if ( + event.affectsConfiguration(SETTING_MCP_PORT) || + event.affectsConfiguration(SETTING_MCP_ENABLED) || + event.affectsConfiguration(SETTING_SHELL_PATH) + ) { + didChangeEmitter.fire(); + } + }), + ); + + context.subscriptions.push(didChangeEmitter); + } catch (err) { + ext.outputChannel.appendLine('error while registering MCP server: ' + String(err)); + } +} diff --git a/src/cosmosDBShell/nodeCredentials.test.ts b/src/cosmosDBShell/nodeCredentials.test.ts new file mode 100644 index 000000000..48c93ac73 --- /dev/null +++ b/src/cosmosDBShell/nodeCredentials.test.ts @@ -0,0 +1,153 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { describe, expect, it, vi } from 'vitest'; +import { AuthenticationMethod } from '../cosmosdb/AuthenticationMethod'; +import { type CosmosDBCredential } from '../cosmosdb/CosmosDBCredential'; +import { type NoSqlContainerResourceItem } from '../tree/nosql/NoSqlContainerResourceItem'; +import { + getCosmosDBShellCredential, + getEntraIdCredential, + getManagedIdentityCredential, + getNodeAuthKind, +} from './nodeCredentials'; + +vi.mock('../extensionVariables', () => ({ + ext: { + outputChannel: { + appendLine: vi.fn(), + }, + }, +})); + +// azureSessionHelper transitively requires @microsoft/vscode-azext-azureauth, which uses +// CJS `require('vscode')` that bypasses the vitest alias for `vscode`. Mock it out — these +// tests only exercise the synchronous credential getters / classifier, not the token path. +vi.mock('../cosmosdb/utils/azureSessionHelper', () => ({ + getAccessTokenForVSCode: vi.fn(), +})); + +type MakeNodeOptions = { + endpoint?: string; + isEmulator?: boolean; + credentials?: CosmosDBCredential[]; +}; + +function makeNode(opts: MakeNodeOptions = {}): NoSqlContainerResourceItem { + return { + model: { + accountInfo: { + endpoint: opts.endpoint ?? 'https://acct.documents.azure.com:443/', + isEmulator: opts.isEmulator ?? false, + credentials: opts.credentials ?? [], + }, + }, + } as unknown as NoSqlContainerResourceItem; +} + +describe('nodeCredentials.getNodeAuthKind', () => { + it('returns "emulator" when isEmulator is true, regardless of credentials', () => { + expect(getNodeAuthKind(makeNode({ isEmulator: true }))).toBe('emulator'); + expect( + getNodeAuthKind( + makeNode({ + isEmulator: true, + credentials: [{ type: AuthenticationMethod.accountKey, key: 'k' }], + }), + ), + ).toBe('emulator'); + }); + + it('returns "accountKey" when an account-key credential exists', () => { + expect(getNodeAuthKind(makeNode({ credentials: [{ type: AuthenticationMethod.accountKey, key: 'k' }] }))).toBe( + 'accountKey', + ); + }); + + it('prefers accountKey over entraId when both are present', () => { + expect( + getNodeAuthKind( + makeNode({ + credentials: [ + { type: AuthenticationMethod.entraId, tenantId: 't' }, + { type: AuthenticationMethod.accountKey, key: 'k' }, + ], + }), + ), + ).toBe('accountKey'); + }); + + it('returns "entraId" when only an entra credential exists', () => { + expect( + getNodeAuthKind(makeNode({ credentials: [{ type: AuthenticationMethod.entraId, tenantId: 't' }] })), + ).toBe('entraId'); + }); + + it('prefers entraId over managedIdentity when both are present (and no account key)', () => { + expect( + getNodeAuthKind( + makeNode({ + credentials: [ + { type: AuthenticationMethod.managedIdentity, clientId: 'c' }, + { type: AuthenticationMethod.entraId, tenantId: 't' }, + ], + }), + ), + ).toBe('entraId'); + }); + + it('returns "managedIdentity" when only a managed-identity credential exists', () => { + expect( + getNodeAuthKind( + makeNode({ + credentials: [{ type: AuthenticationMethod.managedIdentity, clientId: 'c' }], + }), + ), + ).toBe('managedIdentity'); + }); + + it('returns "none" when the credentials list is empty', () => { + expect(getNodeAuthKind(makeNode({}))).toBe('none'); + }); +}); + +describe('nodeCredentials credential getters', () => { + it('getCosmosDBShellCredential returns the key for the first account-key credential', () => { + expect( + getCosmosDBShellCredential( + makeNode({ credentials: [{ type: AuthenticationMethod.accountKey, key: 'KEY' }] }), + ), + ).toBe('KEY'); + }); + + it('getCosmosDBShellCredential returns undefined when no account-key credential is present', () => { + expect( + getCosmosDBShellCredential( + makeNode({ credentials: [{ type: AuthenticationMethod.entraId, tenantId: 't' }] }), + ), + ).toBeUndefined(); + }); + + it('getEntraIdCredential returns the entra credential object', () => { + const cred: CosmosDBCredential = { type: AuthenticationMethod.entraId, tenantId: 'TENANT' }; + expect(getEntraIdCredential(makeNode({ credentials: [cred] }))).toBe(cred); + }); + + it('getEntraIdCredential returns undefined when no entra credential is present', () => { + expect(getEntraIdCredential(makeNode({}))).toBeUndefined(); + }); + + it('getManagedIdentityCredential returns the managed-identity credential object', () => { + const cred: CosmosDBCredential = { + type: AuthenticationMethod.managedIdentity, + clientId: 'CID', + }; + expect(getManagedIdentityCredential(makeNode({ credentials: [cred] }))).toBe(cred); + }); + + it('getManagedIdentityCredential returns undefined when no managed-identity credential is present', () => { + expect(getManagedIdentityCredential(makeNode({}))).toBeUndefined(); + }); +}); diff --git a/src/cosmosDBShell/nodeCredentials.ts b/src/cosmosDBShell/nodeCredentials.ts new file mode 100644 index 000000000..6a457b7cb --- /dev/null +++ b/src/cosmosDBShell/nodeCredentials.ts @@ -0,0 +1,72 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +/** + * Credential extraction and authentication classification for {@link NoSqlContainerResourceItem} + * nodes. The {@link AuthKind} value is what the terminal-reuse layer consults to decide whether + * an already-running shell can host a new node. + */ +import { AuthenticationMethod } from '../cosmosdb/AuthenticationMethod'; +import { type CosmosDBEntraIdCredential, type CosmosDBManagedIdentityCredential } from '../cosmosdb/CosmosDBCredential'; +import { getAccessTokenForVSCode } from '../cosmosdb/utils/azureSessionHelper'; +import { ext } from '../extensionVariables'; +import { type NoSqlContainerResourceItem } from '../tree/nosql/NoSqlContainerResourceItem'; + +/** Authentication mode used at launch — determines which env vars (if any) are baked into the process. */ +export type AuthKind = 'emulator' | 'accountKey' | 'entraId' | 'managedIdentity' | 'none'; + +export function getCosmosDBShellCredential(node: NoSqlContainerResourceItem): string | undefined { + const credential = node.model.accountInfo.credentials.find((c) => c.type === AuthenticationMethod.accountKey); + return credential?.key; +} + +export function getEntraIdCredential(node: NoSqlContainerResourceItem): CosmosDBEntraIdCredential | undefined { + return node.model.accountInfo.credentials.find((c) => c.type === AuthenticationMethod.entraId); +} + +export function getManagedIdentityCredential( + node: NoSqlContainerResourceItem, +): CosmosDBManagedIdentityCredential | undefined { + return node.model.accountInfo.credentials.find((c) => c.type === AuthenticationMethod.managedIdentity); +} + +/** + * Classifies the authentication mode required by a node. This determines which env vars + * (if any) the shell process must have been launched with in order to authenticate. + */ +export function getNodeAuthKind(node: NoSqlContainerResourceItem): AuthKind { + if (node.model.accountInfo.isEmulator) { + return 'emulator'; + } + if (getCosmosDBShellCredential(node)) { + return 'accountKey'; + } + if (getEntraIdCredential(node)) { + return 'entraId'; + } + if (getManagedIdentityCredential(node)) { + return 'managedIdentity'; + } + return 'none'; +} + +/** + * Obtains an access token from VS Code's authentication session for the Cosmos DB endpoint. + * Used as a fallback token via COSMOSDB_SHELL_TOKEN if VisualStudioCodeCredential fails in the shell. + */ +export async function getCosmosDBShellToken( + entraCredential: CosmosDBEntraIdCredential, + endpoint: string, +): Promise { + try { + const endpointUrl = new URL(endpoint); + const scope = `${endpointUrl.origin}${endpointUrl.pathname}.default`; + const token = await getAccessTokenForVSCode(scope, entraCredential.tenantId, { createIfNone: false }); + return token?.token ?? undefined; + } catch { + ext.outputChannel.appendLine('Failed to obtain fallback access token for Cosmos DB Shell'); + return undefined; + } +} diff --git a/src/cosmosDBShell/shellCommand.test.ts b/src/cosmosDBShell/shellCommand.test.ts new file mode 100644 index 000000000..aec703d45 --- /dev/null +++ b/src/cosmosDBShell/shellCommand.test.ts @@ -0,0 +1,127 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import * as fs from 'fs'; +import { beforeEach, describe, expect, it, vi, type Mock } from 'vitest'; +import { SettingsService } from '../services/SettingsService'; +import { isCosmosDBShellPathFound, quoteArg } from './shellCommand'; + +vi.mock('fs', () => ({ + existsSync: vi.fn(), + statSync: vi.fn(), +})); + +vi.mock('../services/SettingsService', () => ({ + SettingsService: { + getSetting: vi.fn(), + }, +})); + +vi.mock('../extensionVariables', () => ({ + ext: { + outputChannel: { + appendLine: vi.fn(), + error: vi.fn(), + }, + }, +})); + +describe('shellCommand.quoteArg', () => { + it('returns the value unchanged when no whitespace, quotes, or backslashes are present', () => { + expect(quoteArg('foo')).toBe('foo'); + expect(quoteArg('connect')).toBe('connect'); + }); + + it('wraps values containing spaces in double quotes', () => { + expect(quoteArg('foo bar')).toBe('"foo bar"'); + expect(quoteArg('https://my account.documents.azure.com/')).toBe('"https://my account.documents.azure.com/"'); + }); + + it('wraps and escapes embedded double quotes', () => { + expect(quoteArg('he said "hi"')).toBe('"he said \\"hi\\""'); + }); + + it('wraps values containing single quotes', () => { + expect(quoteArg("don't")).toBe('"don\'t"'); + }); + + it('does not strip leading or trailing whitespace', () => { + expect(quoteArg(' foo ')).toBe('" foo "'); + }); + + it('wraps and escapes embedded backslashes', () => { + // Without backslash escaping, `foo\"bar` would round-trip to `"foo\"bar"`, where + // the `\"` is parsed by the shell as an escaped quote and the actual `"` then + // terminates the argument early. Backslashes must be escaped before quotes. + expect(quoteArg('foo\\bar')).toBe('"foo\\\\bar"'); + expect(quoteArg('foo\\"bar')).toBe('"foo\\\\\\"bar"'); + }); + + it('wraps and escapes Windows-style paths containing backslashes', () => { + expect(quoteArg('C:\\Users\\test')).toBe('"C:\\\\Users\\\\test"'); + expect(quoteArg('C:\\path with space\\file.txt')).toBe('"C:\\\\path with space\\\\file.txt"'); + }); +}); + +describe('shellCommand.isCosmosDBShellPathFound', () => { + beforeEach(() => { + vi.clearAllMocks(); + (fs.existsSync as Mock).mockReturnValue(false); + (fs.statSync as Mock).mockReturnValue({ isFile: () => false }); + }); + + it('returns false when no path is configured', () => { + (SettingsService.getSetting as Mock).mockReturnValue(undefined); + expect(isCosmosDBShellPathFound()).toBe(false); + }); + + it('returns false when the configured path is whitespace only', () => { + (SettingsService.getSetting as Mock).mockReturnValue(' '); + expect(isCosmosDBShellPathFound()).toBe(false); + }); + + it('strips wrapping double quotes before checking the filesystem', () => { + const path = 'C:\\tools\\shell.exe'; + (SettingsService.getSetting as Mock).mockReturnValue(`"${path}"`); + (fs.existsSync as Mock).mockImplementation((p: string) => p === path); + (fs.statSync as Mock).mockReturnValue({ isFile: () => true }); + + expect(isCosmosDBShellPathFound()).toBe(true); + }); + + it('strips wrapping single quotes before checking the filesystem', () => { + const path = '/usr/local/bin/cosmosdbshell'; + (SettingsService.getSetting as Mock).mockReturnValue(`'${path}'`); + (fs.existsSync as Mock).mockImplementation((p: string) => p === path); + (fs.statSync as Mock).mockReturnValue({ isFile: () => true }); + + expect(isCosmosDBShellPathFound()).toBe(true); + }); + + it('returns false when the configured path does not exist on disk', () => { + (SettingsService.getSetting as Mock).mockReturnValue('C:\\does\\not\\exist.exe'); + (fs.existsSync as Mock).mockReturnValue(false); + + expect(isCosmosDBShellPathFound()).toBe(false); + }); + + it('returns false when the configured path resolves to a directory', () => { + (SettingsService.getSetting as Mock).mockReturnValue('/usr/bin'); + (fs.existsSync as Mock).mockReturnValue(true); + (fs.statSync as Mock).mockReturnValue({ isFile: () => false }); + + expect(isCosmosDBShellPathFound()).toBe(false); + }); + + it('returns false when statSync throws (e.g. permission denied)', () => { + (SettingsService.getSetting as Mock).mockReturnValue('/no/access'); + (fs.existsSync as Mock).mockReturnValue(true); + (fs.statSync as Mock).mockImplementation(() => { + throw new Error('EACCES'); + }); + + expect(isCosmosDBShellPathFound()).toBe(false); + }); +}); diff --git a/src/cosmosDBShell/shellCommand.ts b/src/cosmosDBShell/shellCommand.ts new file mode 100644 index 000000000..cd7bef801 --- /dev/null +++ b/src/cosmosDBShell/shellCommand.ts @@ -0,0 +1,89 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +/** + * Low-level primitives for locating and launching the CosmosDBShell binary. + * Kept dependency-light so they can be reused by every other module in this folder. + */ +import * as fs from 'fs'; +import * as vscode from 'vscode'; +import { ext } from '../extensionVariables'; +import { SettingsService } from '../services/SettingsService'; +import { SETTING_SHELL_PATH } from './constants'; +import { resolveCosmosDBShellCommand } from './cosmosDBShellCommandResolver'; + +/** + * Resolves the user-configured Cosmos DB Shell command (or the default `cosmosdbshell`) + * to a runnable executable path. Delegates Windows PATH/.cmd shim resolution to + * {@link resolveCosmosDBShellCommand}. + */ +export function getCosmosDBShellCommand(): string { + const shellPath: string | undefined = SettingsService.getSetting(SETTING_SHELL_PATH); + return resolveCosmosDBShellCommand(shellPath); +} + +/** + * Returns true when the user has configured a shell path and that path points at an + * existing file on disk. Used to differentiate "missing binary" from "misconfigured path" + * when surfacing the install/repair prompt. + */ +export function isCosmosDBShellPathFound(): boolean { + const shellPath: string | undefined = SettingsService.getSetting(SETTING_SHELL_PATH); + if (!shellPath?.trim()) { + return false; + } + + const trimmed = shellPath.trim(); + const unquoted = + (trimmed.startsWith('"') && trimmed.endsWith('"')) || (trimmed.startsWith("'") && trimmed.endsWith("'")) + ? trimmed.slice(1, -1) + : trimmed; + + try { + return fs.existsSync(unquoted) && fs.statSync(unquoted).isFile(); + } catch { + return false; + } +} + +/** + * Watches for the terminal closing shortly after creation (early exit). + * If the process exits quickly, logs the exit code and reason to the output channel. + */ +export function watchForEarlyExit(terminal: vscode.Terminal): void { + const startTime = Date.now(); + const listener = vscode.window.onDidCloseTerminal((closedTerminal) => { + if (closedTerminal !== terminal) { + return; + } + + clearTimeout(timeout); + listener.dispose(); + if (Date.now() - startTime < 5000) { + const exitCode = closedTerminal.exitStatus?.code; + const exitReason = closedTerminal.exitStatus?.reason; + ext.outputChannel.error( + `Cosmos DB Shell exited early.${exitCode !== undefined ? ` Exit code: ${exitCode}.` : ''}${exitReason !== undefined ? ` Reason: ${exitReason}.` : ''}`, + ); + } + }); + const timeout = setTimeout(() => { + listener.dispose(); + }, 5000); +} + +/** + * Quotes a value for use as a single argument inside an interactive Cosmos DB Shell + * command (sent via `terminal.sendText`). Wraps values containing whitespace, quotes, or + * backslashes in double quotes and applies C-style escapes (`\\` and `\"`) so the + * argument is preserved verbatim when parsed by the shell. + * + * Backslashes are escaped before quotes so the substitutions don't compound (e.g. a + * literal `\"` in the input becomes `\\\"` in the output, not `\\"` which would be + * parsed as `\` followed by an argument-terminating quote). + */ +export function quoteArg(value: string): string { + return /[\s"'\\]/.test(value) ? `"${value.replace(/\\/g, '\\\\').replace(/"/g, '\\"')}"` : value; +} diff --git a/src/cosmosDBShell/shellSupportCache.test.ts b/src/cosmosDBShell/shellSupportCache.test.ts new file mode 100644 index 000000000..df4561f90 --- /dev/null +++ b/src/cosmosDBShell/shellSupportCache.test.ts @@ -0,0 +1,128 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import * as child from 'child_process'; +import { beforeEach, describe, expect, it, vi, type Mock } from 'vitest'; +import { getCosmosDBShellCommand } from './shellCommand'; +import { + getDetectedCosmosDBShellVersion, + invalidateCosmosDBShellSupportCache, + isCosmosDBShellInstalled, +} from './shellSupportCache'; + +vi.mock('child_process', () => ({ + execFileSync: vi.fn(), +})); + +vi.mock('./shellCommand', () => ({ + getCosmosDBShellCommand: vi.fn(), +})); + +vi.mock('../extensionVariables', () => ({ + ext: { + outputChannel: { + appendLine: vi.fn(), + }, + }, +})); + +describe('shellSupportCache', () => { + beforeEach(() => { + invalidateCosmosDBShellSupportCache(); + vi.clearAllMocks(); + (getCosmosDBShellCommand as Mock).mockReturnValue('cosmosdbshell'); + }); + + it('marks the shell as installed and parses the version on a clean successful exit', () => { + (child.execFileSync as Mock).mockReturnValue(Buffer.from('CosmosDBShell 1.2.3\n')); + expect(isCosmosDBShellInstalled()).toBe(true); + expect(getDetectedCosmosDBShellVersion()).toBe('1.2.3'); + }); + + it('parses pre-release version tokens (e.g. 1.2.3-prerelease.4)', () => { + (child.execFileSync as Mock).mockReturnValue(Buffer.from('CosmosDBShell 1.2.3-prerelease.4\n')); + expect(getDetectedCosmosDBShellVersion()).toBe('1.2.3-prerelease.4'); + }); + + it('returns installed=true with undefined version when no SemVer-like token is in the output', () => { + (child.execFileSync as Mock).mockReturnValue(Buffer.from('no version info here\n')); + expect(isCosmosDBShellInstalled()).toBe(true); + expect(getDetectedCosmosDBShellVersion()).toBeUndefined(); + }); + + it('caches the probe result across calls keyed on the resolved command', () => { + (child.execFileSync as Mock).mockReturnValue(Buffer.from('1.0.0\n')); + isCosmosDBShellInstalled(); + isCosmosDBShellInstalled(); + getDetectedCosmosDBShellVersion(); + expect((child.execFileSync as Mock).mock.calls).toHaveLength(1); + }); + + it('invalidate clears the cache so the next call re-probes', () => { + (child.execFileSync as Mock).mockReturnValue(Buffer.from('1.0.0\n')); + isCosmosDBShellInstalled(); + invalidateCosmosDBShellSupportCache(); + isCosmosDBShellInstalled(); + expect((child.execFileSync as Mock).mock.calls).toHaveLength(2); + }); + + it('uses a separate cache entry per resolved command path', () => { + (child.execFileSync as Mock).mockReturnValue(Buffer.from('1.0.0\n')); + + (getCosmosDBShellCommand as Mock).mockReturnValue('cmd-a'); + isCosmosDBShellInstalled(); + + (getCosmosDBShellCommand as Mock).mockReturnValue('cmd-b'); + isCosmosDBShellInstalled(); + + expect((child.execFileSync as Mock).mock.calls).toHaveLength(2); + }); + + it('treats non-zero exit with recognizable shell output as installed (ANSI fallback)', () => { + const err = Object.assign(new Error('Spawn failed'), { + stdout: Buffer.from('CosmosDBShell 2.0.0-rc.1\n'), + stderr: Buffer.from('ansi: terminal not supported'), + }); + (child.execFileSync as Mock).mockImplementation(() => { + throw err; + }); + + expect(isCosmosDBShellInstalled()).toBe(true); + expect(getDetectedCosmosDBShellVersion()).toBe('2.0.0-rc.1'); + }); + + it('handles string-typed stdout/stderr on the error object (not only Buffers)', () => { + const err = Object.assign(new Error('Spawn failed'), { + stdout: 'CosmosDBShell 3.0.0\n', + stderr: '', + }); + (child.execFileSync as Mock).mockImplementation(() => { + throw err; + }); + + expect(isCosmosDBShellInstalled()).toBe(true); + expect(getDetectedCosmosDBShellVersion()).toBe('3.0.0'); + }); + + it('treats non-zero exit with no recognizable output as not installed', () => { + const err = Object.assign(new Error('spawn ENOENT'), { + stdout: '', + stderr: '', + }); + (child.execFileSync as Mock).mockImplementation(() => { + throw err; + }); + + expect(isCosmosDBShellInstalled()).toBe(false); + expect(getDetectedCosmosDBShellVersion()).toBeUndefined(); + }); + + it('treats a thrown error with no stdout/stderr at all as not installed', () => { + (child.execFileSync as Mock).mockImplementation(() => { + throw new Error('boom'); + }); + expect(isCosmosDBShellInstalled()).toBe(false); + }); +}); diff --git a/src/cosmosDBShell/shellSupportCache.ts b/src/cosmosDBShell/shellSupportCache.ts new file mode 100644 index 000000000..71a1a138c --- /dev/null +++ b/src/cosmosDBShell/shellSupportCache.ts @@ -0,0 +1,122 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +/** + * Cached probing of the CosmosDBShell binary via `--version`. + * + * Caching avoids re-spawning the shell on every keystroke (the value is consulted by + * activation, by the MCP provider's discovery filter, by the language-server bootstrap, + * etc.). The cache is keyed on the resolved command path and is invalidated whenever + * the `cosmosDB.shell.path` setting changes. + */ +import * as child from 'child_process'; +import { ext } from '../extensionVariables'; +import { getCosmosDBShellCommand } from './shellCommand'; + +type CosmosDBShellSupportInfo = { installed: boolean; version?: string }; + +const cosmosDBShellSupportCache = new Map(); + +/** + * Determines if CosmosDBShell is installed. + * + * @returns true, if CosmosDBShell is installed, false otherwise. + */ +export function isCosmosDBShellInstalled(): boolean { + return getCachedShellSupport().installed; +} + +/** + * Returns the version reported by `CosmosDBShell --version` (e.g. `1.2.3` or + * `1.2.3-prerelease.45`), or undefined when the shell is not installed or no + * version could be parsed from its output. + */ +export function getDetectedCosmosDBShellVersion(): string | undefined { + return getCachedShellSupport().version; +} + +/** + * Clears the cached result of {@link isCosmosDBShellInstalled}. + * Call this when the shell path configuration changes or the binary may have been installed/removed. + */ +export function invalidateCosmosDBShellSupportCache(): void { + cosmosDBShellSupportCache.clear(); +} + +function getCachedShellSupport(): CosmosDBShellSupportInfo { + const command = getCosmosDBShellCommand(); + const cached = cosmosDBShellSupportCache.get(command); + if (cached !== undefined) { + return cached; + } + const result = detectCosmosDBShellSupport(command); + cosmosDBShellSupportCache.set(command, result); + return result; +} + +/** + * Extracts a SemVer-like version token (e.g. `1.2.3` or `1.2.3-prerelease.4`) + * from the `--version` output of CosmosDBShell. Returns undefined when no + * recognizable version token is present. + */ +function parseShellVersion(output: string): string | undefined { + const match = output.match(/\b(\d+\.\d+\.\d+(?:[-+][0-9A-Za-z.+-]+)?)\b/); + return match?.[1]; +} + +function detectCosmosDBShellSupport(command: string): CosmosDBShellSupportInfo { + try { + const stdout = child.execFileSync(command, ['--version'], { + windowsHide: true, + env: { + ...process.env, + // CosmosDBShell may use ANSI output libraries (e.g. Spectre.Console). + // When spawned by VS Code, stdio is typically redirected which can confuse + // terminal capability detection. Prefer a plain output mode. + NO_COLOR: '1', + CLICOLOR: '0', + TERM: process.env.TERM ?? 'dumb', + }, + }); + return { installed: true, version: parseShellVersion(stdout.toString('utf8')) }; + } catch (err) { + const anyErr = err as { stdout?: unknown; stderr?: unknown }; + const stdout = + typeof anyErr?.stdout === 'string' + ? anyErr.stdout + : Buffer.isBuffer(anyErr?.stdout) + ? anyErr.stdout.toString('utf8') + : ''; + const stderr = + typeof anyErr?.stderr === 'string' + ? anyErr.stderr + : Buffer.isBuffer(anyErr?.stderr) + ? anyErr.stderr.toString('utf8') + : ''; + + // Workaround: CosmosDBShell may print a valid version string but still exit non-zero + // when ANSI is not available. Treat that as installed. + const combinedOutput = `${stdout}\n${stderr}`; + if (/\bCosmos(?:DB)?Shell\b/i.test(combinedOutput)) { + ext.outputChannel.appendLine( + 'warning: CosmosDBShell "--version" exited non-zero, but returned version output; treating as installed.', + ); + if (stderr.trim().length > 0) { + ext.outputChannel.appendLine(stderr.trim()); + } + return { installed: true, version: parseShellVersion(combinedOutput) }; + } + + ext.outputChannel.appendLine('fail ' + String(err)); + ext.outputChannel.appendLine('while running "' + command + ' --version"'); + if (stdout.trim().length > 0) { + ext.outputChannel.appendLine('stdout: ' + stdout.trim()); + } + if (stderr.trim().length > 0) { + ext.outputChannel.appendLine('stderr: ' + stderr.trim()); + } + return { installed: false }; + } +} diff --git a/src/cosmosDBShell/terminalReuse.test.ts b/src/cosmosDBShell/terminalReuse.test.ts new file mode 100644 index 000000000..62188b7a1 --- /dev/null +++ b/src/cosmosDBShell/terminalReuse.test.ts @@ -0,0 +1,243 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { describe, expect, it, vi } from 'vitest'; +import { AuthenticationMethod } from '../cosmosdb/AuthenticationMethod'; +import { type CosmosDBCredential } from '../cosmosdb/CosmosDBCredential'; +import { type NoSqlContainerResourceItem } from '../tree/nosql/NoSqlContainerResourceItem'; +import { + buildInteractiveConnectCommand, + buildTerminalStateForNode, + canReuseTerminalForNode, + type ShellTerminalState, +} from './terminalReuse'; + +vi.mock('../extensionVariables', () => ({ + ext: { + outputChannel: { + appendLine: vi.fn(), + }, + }, +})); + +// nodeCredentials transitively loads @microsoft/vscode-azext-azureauth, which uses CJS +// `require('vscode')` that bypasses the vitest alias. Stub it for module load. +vi.mock('../cosmosdb/utils/azureSessionHelper', () => ({ + getAccessTokenForVSCode: vi.fn(), +})); + +type MakeNodeOptions = { + endpoint?: string; + isEmulator?: boolean; + credentials?: CosmosDBCredential[]; +}; + +function makeNode(opts: MakeNodeOptions = {}): NoSqlContainerResourceItem { + return { + model: { + accountInfo: { + endpoint: opts.endpoint ?? 'https://acct.documents.azure.com:443/', + isEmulator: opts.isEmulator ?? false, + credentials: opts.credentials ?? [], + }, + }, + } as unknown as NoSqlContainerResourceItem; +} + +describe('terminalReuse.buildTerminalStateForNode', () => { + it('captures endpoint, auth kind, tenant, and managed-identity client id', () => { + const node = makeNode({ + endpoint: 'https://x/', + credentials: [ + { type: AuthenticationMethod.entraId, tenantId: 'T1' }, + { type: AuthenticationMethod.managedIdentity, clientId: 'MI1' }, + ], + }); + expect(buildTerminalStateForNode(node)).toEqual({ + endpoint: 'https://x/', + authKind: 'entraId', + tenantId: 'T1', + managedIdentityClientId: 'MI1', + }); + }); + + it('defaults endpoint to "" when accountInfo.endpoint is undefined', () => { + const node = { + model: { accountInfo: { credentials: [], isEmulator: false } }, + } as unknown as NoSqlContainerResourceItem; + expect(buildTerminalStateForNode(node).endpoint).toBe(''); + }); + + it('reports authKind = "accountKey" when an account key is present', () => { + const node = makeNode({ + credentials: [{ type: AuthenticationMethod.accountKey, key: 'k' }], + }); + expect(buildTerminalStateForNode(node).authKind).toBe('accountKey'); + }); + + it('reports authKind = "emulator" when the node is an emulator', () => { + expect(buildTerminalStateForNode(makeNode({ isEmulator: true })).authKind).toBe('emulator'); + }); + + it('reports authKind = "none" when no credentials are present', () => { + const state = buildTerminalStateForNode(makeNode({})); + expect(state.authKind).toBe('none'); + expect(state.tenantId).toBeUndefined(); + expect(state.managedIdentityClientId).toBeUndefined(); + }); +}); + +describe('terminalReuse.canReuseTerminalForNode', () => { + const accountKeyTerminal: ShellTerminalState = { + endpoint: 'https://acct/', + authKind: 'accountKey', + }; + + it('allows reuse for emulator nodes regardless of terminal state', () => { + expect(canReuseTerminalForNode(accountKeyTerminal, makeNode({ isEmulator: true }))).toBe(true); + }); + + it('allows reuse for managed-identity nodes regardless of terminal state', () => { + expect( + canReuseTerminalForNode( + accountKeyTerminal, + makeNode({ + credentials: [{ type: AuthenticationMethod.managedIdentity, clientId: 'mi' }], + }), + ), + ).toBe(true); + }); + + it('allows reuse for "none" auth nodes regardless of terminal state', () => { + expect(canReuseTerminalForNode(accountKeyTerminal, makeNode({}))).toBe(true); + }); + + it('rejects reuse when account-key node endpoint differs from terminal endpoint', () => { + expect( + canReuseTerminalForNode( + { endpoint: 'https://a/', authKind: 'accountKey' }, + makeNode({ + endpoint: 'https://b/', + credentials: [{ type: AuthenticationMethod.accountKey, key: 'k' }], + }), + ), + ).toBe(false); + }); + + it('rejects reuse when account-key node matches endpoint but terminal authKind differs', () => { + expect( + canReuseTerminalForNode( + { endpoint: 'https://a/', authKind: 'entraId' }, + makeNode({ + endpoint: 'https://a/', + credentials: [{ type: AuthenticationMethod.accountKey, key: 'k' }], + }), + ), + ).toBe(false); + }); + + it('allows reuse for account-key node when endpoint and authKind match', () => { + expect( + canReuseTerminalForNode( + { endpoint: 'https://a/', authKind: 'accountKey' }, + makeNode({ + endpoint: 'https://a/', + credentials: [{ type: AuthenticationMethod.accountKey, key: 'k' }], + }), + ), + ).toBe(true); + }); + + it('rejects reuse for entra node when tenant differs from terminal tenant', () => { + expect( + canReuseTerminalForNode( + { endpoint: 'https://a/', authKind: 'entraId', tenantId: 'T1' }, + makeNode({ + endpoint: 'https://a/', + credentials: [{ type: AuthenticationMethod.entraId, tenantId: 'T2' }], + }), + ), + ).toBe(false); + }); + + it('allows reuse for entra node when endpoint and tenant match', () => { + expect( + canReuseTerminalForNode( + { endpoint: 'https://a/', authKind: 'entraId', tenantId: 'T1' }, + makeNode({ + endpoint: 'https://a/', + credentials: [{ type: AuthenticationMethod.entraId, tenantId: 'T1' }], + }), + ), + ).toBe(true); + }); +}); + +describe('terminalReuse.buildInteractiveConnectCommand', () => { + it('emits a plain connect for emulator nodes (no credential flags)', () => { + const node = makeNode({ + isEmulator: true, + endpoint: 'https://emu:8081/', + credentials: [{ type: AuthenticationMethod.entraId, tenantId: 'T' }], + }); + expect(buildInteractiveConnectCommand(node, 'https://emu:8081/')).toBe('connect https://emu:8081/'); + }); + + it('appends --vscode-credential and --tenant for entra nodes', () => { + const node = makeNode({ + credentials: [{ type: AuthenticationMethod.entraId, tenantId: 'TENANT' }], + }); + expect(buildInteractiveConnectCommand(node, 'https://a/')).toBe( + 'connect https://a/ --vscode-credential --tenant TENANT', + ); + }); + + it('appends --vscode-credential without --tenant when no tenantId is set on the entra credential', () => { + const node = makeNode({ + credentials: [{ type: AuthenticationMethod.entraId, tenantId: undefined }], + }); + expect(buildInteractiveConnectCommand(node, 'https://a/')).toBe('connect https://a/ --vscode-credential'); + }); + + it('appends --managed-identity with the client id for managed-identity nodes', () => { + const node = makeNode({ + credentials: [{ type: AuthenticationMethod.managedIdentity, clientId: 'MI-CID' }], + }); + expect(buildInteractiveConnectCommand(node, 'https://a/')).toBe('connect https://a/ --managed-identity MI-CID'); + }); + + it('omits --managed-identity when the managed-identity credential has no clientId', () => { + const node = makeNode({ + credentials: [{ type: AuthenticationMethod.managedIdentity, clientId: undefined }], + }); + expect(buildInteractiveConnectCommand(node, 'https://a/')).toBe('connect https://a/'); + }); + + it('quotes endpoints containing spaces or quotes', () => { + const node = makeNode({}); + expect(buildInteractiveConnectCommand(node, 'https://a b/')).toBe('connect "https://a b/"'); + }); + + it('quotes tenant ids containing spaces', () => { + const node = makeNode({ + credentials: [{ type: AuthenticationMethod.entraId, tenantId: 'a b' }], + }); + expect(buildInteractiveConnectCommand(node, 'https://a/')).toBe( + 'connect https://a/ --vscode-credential --tenant "a b"', + ); + }); + + it('emits both --vscode-credential and --managed-identity when both credentials are present', () => { + const node = makeNode({ + credentials: [ + { type: AuthenticationMethod.entraId, tenantId: 'T' }, + { type: AuthenticationMethod.managedIdentity, clientId: 'MI' }, + ], + }); + expect(buildInteractiveConnectCommand(node, 'https://a/')).toBe( + 'connect https://a/ --vscode-credential --tenant T --managed-identity MI', + ); + }); +}); diff --git a/src/cosmosDBShell/terminalReuse.ts b/src/cosmosDBShell/terminalReuse.ts new file mode 100644 index 000000000..c50d9c51a --- /dev/null +++ b/src/cosmosDBShell/terminalReuse.ts @@ -0,0 +1,117 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +/** + * Tracks per-terminal launch state so we know which Cosmos DB Shell terminals can be reused + * for a given node. The state describes how the process was *launched* (endpoint + auth + * mode + env vars baked in), not its current in-shell connection status: a user may have + * run `disconnect` inside the shell, which VS Code cannot observe. The reuse path therefore + * always re-issues `connect` before sending further commands. + */ +import * as vscode from 'vscode'; +import { type NoSqlContainerResourceItem } from '../tree/nosql/NoSqlContainerResourceItem'; +import { type AuthKind, getEntraIdCredential, getManagedIdentityCredential, getNodeAuthKind } from './nodeCredentials'; +import { quoteArg } from './shellCommand'; + +export type ShellTerminalState = { + /** Endpoint the shell process was launched against, or '' for command-palette launches without a node. */ + endpoint: string; + /** Authentication mode used at launch. Determines which env vars (if any) are baked into the process. */ + authKind: AuthKind; + tenantId?: string; + managedIdentityClientId?: string; +}; + +export const terminalStates = new Map(); + +/** Builds a {@link ShellTerminalState} record describing how a shell would be launched for this node. */ +export function buildTerminalStateForNode(node: NoSqlContainerResourceItem): ShellTerminalState { + return { + endpoint: node.model.accountInfo.endpoint ?? '', + authKind: getNodeAuthKind(node), + tenantId: getEntraIdCredential(node)?.tenantId, + managedIdentityClientId: getManagedIdentityCredential(node)?.clientId, + }; +} + +/** + * Determines whether an already-running Cosmos DB Shell terminal can host the given node. + * + * Auth modes that need launch-time env vars (account key, Entra ID fallback token) are only + * compatible if the terminal was launched for the *same endpoint* with the *same* auth mode + * (and tenant for Entra ID) — otherwise the baked-in env would be wrong for the new node. + * Auth modes that don't rely on env vars (emulator, managed identity, none) can run in any + * tracked terminal via the interactive `connect` command. + */ +export function canReuseTerminalForNode(state: ShellTerminalState, node: NoSqlContainerResourceItem): boolean { + const nodeAuth = getNodeAuthKind(node); + + if (nodeAuth === 'emulator' || nodeAuth === 'managedIdentity' || nodeAuth === 'none') { + return true; + } + + if (state.endpoint !== node.model.accountInfo.endpoint || state.authKind !== nodeAuth) { + return false; + } + + if (nodeAuth === 'entraId') { + const cred = getEntraIdCredential(node); + if (cred?.tenantId !== state.tenantId) { + return false; + } + } + + return true; +} + +/** + * Finds the best tracked Cosmos DB Shell terminal to reuse for the given node, preferring + * terminals already associated with the same endpoint to keep terminal usage stable. + */ +export function findReusableTerminalForNode( + node: NoSqlContainerResourceItem, +): { terminal: vscode.Terminal; state: ShellTerminalState } | undefined { + const candidates: Array<{ terminal: vscode.Terminal; state: ShellTerminalState; sameEndpoint: boolean }> = []; + for (const [terminal, state] of terminalStates) { + if (!vscode.window.terminals.includes(terminal)) { + continue; + } + if (!canReuseTerminalForNode(state, node)) { + continue; + } + candidates.push({ + terminal, + state, + sameEndpoint: state.endpoint === node.model.accountInfo.endpoint, + }); + } + candidates.sort((a, b) => Number(b.sameEndpoint) - Number(a.sameEndpoint)); + return candidates[0]; +} + +/** + * Builds the interactive `connect` command that mirrors the CLI `--connect` flag and related + * credential flags, so an already-running Cosmos DB Shell can be attached to a specific account. + */ +export function buildInteractiveConnectCommand(node: NoSqlContainerResourceItem, endpoint: string): string { + const parts = ['connect', quoteArg(endpoint)]; + + if (!node.model.accountInfo.isEmulator) { + const entraCredential = getEntraIdCredential(node); + if (entraCredential) { + parts.push('--vscode-credential'); + if (entraCredential.tenantId) { + parts.push('--tenant', quoteArg(entraCredential.tenantId)); + } + } + + const managedIdentityCredential = getManagedIdentityCredential(node); + if (managedIdentityCredential?.clientId) { + parts.push('--managed-identity', quoteArg(managedIdentityCredential.clientId)); + } + } + + return parts.join(' '); +}