diff --git a/extensions/authentication/package.json b/extensions/authentication/package.json index d166de7056f8..222a3159c98e 100644 --- a/extensions/authentication/package.json +++ b/extensions/authentication/package.json @@ -20,6 +20,7 @@ "onAuthenticationRequest:google", "onAuthenticationRequest:google-cloud", "onAuthenticationRequest:deepseek-api", + "onAuthenticationRequest:databricks", "onCommand:authentication.configureProviders", "onCommand:authentication.migrateApiKey" ], @@ -73,6 +74,10 @@ { "id": "deepseek-api", "label": "DeepSeek" + }, + { + "id": "databricks", + "label": "Databricks" } ], "commands": [ @@ -278,6 +283,32 @@ "experimental" ] }, + "authentication.databricks.credentials": { + "type": "object", + "default": {}, + "markdownDescription": "%configuration.databricks.credentials.description%", + "properties": { + "DATABRICKS_HOST": { + "type": "string", + "default": null + } + }, + "additionalProperties": false, + "tags": [ + "experimental" + ] + }, + "authentication.databricks.customHeaders": { + "type": "object", + "default": {}, + "additionalProperties": { + "type": "string" + }, + "markdownDescription": "%configuration.databricks.customHeaders.description%", + "tags": [ + "experimental" + ] + }, "assistant.provider.googleVertex.enabled": { "type": "boolean", "default": true, @@ -294,6 +325,14 @@ "experimental" ] }, + "assistant.provider.databricks.enabled": { + "type": "boolean", + "default": false, + "markdownDescription": "%configuration.provider.databricks.enabled.description%", + "tags": [ + "experimental" + ] + }, "positron.assistant.provider.anthropic.enable": { "type": "boolean", "default": true, diff --git a/extensions/authentication/package.nls.json b/extensions/authentication/package.nls.json index cf7541f2b429..173efcebb050 100644 --- a/extensions/authentication/package.nls.json +++ b/extensions/authentication/package.nls.json @@ -18,6 +18,9 @@ "configuration.openai-compatible.customHeaders.description": "Extra HTTP headers attached to every OpenAI-compatible request. Useful for enterprise gateways that require tenancy or routing headers. Entries may override internal headers and break the connection.", "configuration.deepseek-api.baseUrl.description": "Custom DeepSeek API base URL. Overrides the default https://api.deepseek.com endpoint. Also set automatically from the DEEPSEEK_BASE_URL environment variable.", "configuration.deepseek-api.customHeaders.description": "Extra HTTP headers attached to every DeepSeek request. Useful for enterprise gateways that require tenancy or routing headers. Entries may override internal headers and break the connection.", + "configuration.databricks.credentials.description": "Variables used to configure the Databricks provider.\n\nExample: to set the Databricks workspace, add an item with key `DATABRICKS_HOST` and the workspace URL as the value.", + "configuration.databricks.customHeaders.description": "Extra HTTP headers attached to every Databricks request. Useful for enterprise gateways that require tenancy or routing headers. Entries may override internal headers and break the connection.", + "configuration.provider.databricks.enabled.description": "Enable Databricks as a language model provider for Positron's AI features.\n\nAuthenticate to the provider in the [language model provider dialog](command:authentication.configureProviders?%5B%7B%22preselectedProviderId%22%3A%22databricks%22%7D%5D).", "configuration.provider.googleVertex.enabled.description": "Enable Gemini Enterprise Agent Platform (formerly Google Vertex AI) as a language model provider for Positron's AI features.\n\nRestart Positron for the changes to take effect. Then, authenticate to the provider in the [language model provider dialog](command:authentication.configureProviders?%5B%7B%22preselectedProviderId%22%3A%22google-cloud%22%7D%5D).", "configuration.provider.deepseek.enabled.description": "Enable DeepSeek as a language model provider for Positron's AI features.\n\nAuthenticate to the provider in the [language model provider dialog](command:authentication.configureProviders?%5B%7B%22preselectedProviderId%22%3A%22deepseek-api%22%7D%5D).", "configuration.provider.anthropic.enable.description": "Enable Anthropic as a language model provider for Positron's AI features.\n\nRestart Positron for the changes to take effect. Then, authenticate to the provider in the [language model provider dialog](command:authentication.configureProviders?%5B%7B%22preselectedProviderId%22%3A%22anthropic-api%22%7D%5D).", diff --git a/extensions/authentication/src/authProvider.ts b/extensions/authentication/src/authProvider.ts index fe229cba1386..783b87e01eef 100644 --- a/extensions/authentication/src/authProvider.ts +++ b/extensions/authentication/src/authProvider.ts @@ -84,6 +84,13 @@ export class AuthProvider return !!this.credentialChain?.preventSignOut; } + /** + * Cancel an in-flight interactive sign-in, if the provider supports + * one (e.g. OAuth providers). The config dialog calls this when the + * user cancels. + */ + cancelSignIn?(): void; + /** Expose session-change events. */ fireSessionsChanged( event: vscode.AuthenticationProviderAuthenticationSessionsChangeEvent diff --git a/extensions/authentication/src/configDialog.ts b/extensions/authentication/src/configDialog.ts index 0434bb8b0096..3082b9670d8d 100644 --- a/extensions/authentication/src/configDialog.ts +++ b/extensions/authentication/src/configDialog.ts @@ -7,7 +7,6 @@ import * as vscode from 'vscode'; import * as positron from 'positron'; import { randomUUID } from 'crypto'; import { AuthProvider } from './authProvider'; -import { PositOAuthProvider } from './positOAuthProvider'; import { FOUNDRY_AUTH_PROVIDER_ID } from './constants'; import { log } from './log'; import { FOUNDRY_MANAGED_CREDENTIALS, SNOWFLAKE_MANAGED_CREDENTIALS, hasManagedCredentials } from './managedCredentials'; @@ -175,9 +174,7 @@ export async function providerAction( } case 'cancel': { const provider = authProviders.get(providerId); - if (provider instanceof PositOAuthProvider) { - provider.cancelSignIn(); - } + provider?.cancelSignIn?.(); break; } default: diff --git a/extensions/authentication/src/constants.ts b/extensions/authentication/src/constants.ts index 164266662587..a51d28710762 100644 --- a/extensions/authentication/src/constants.ts +++ b/extensions/authentication/src/constants.ts @@ -22,3 +22,9 @@ export const CUSTOM_PROVIDER_AUTH_PROVIDER_ID = 'openai-compatible'; export const GEMINI_AUTH_PROVIDER_ID = 'google'; export const GOOGLE_CLOUD_AUTH_PROVIDER_ID = 'google-cloud'; export const DEEPSEEK_AUTH_PROVIDER_ID = 'deepseek-api'; +export const DATABRICKS_AUTH_PROVIDER_ID = 'databricks'; + +export const DATABRICKS_OAUTH_CLIENT_ID = 'databricks-cli'; +export const DATABRICKS_OAUTH_REDIRECT_PORT = 8020; +export const DATABRICKS_OAUTH_REDIRECT_URI = 'http://localhost:8020'; +export const DATABRICKS_OAUTH_SCOPES = 'all-apis offline_access'; diff --git a/extensions/authentication/src/databricksAuthProvider.ts b/extensions/authentication/src/databricksAuthProvider.ts new file mode 100644 index 000000000000..cf529feac2d3 --- /dev/null +++ b/extensions/authentication/src/databricksAuthProvider.ts @@ -0,0 +1,347 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (C) 2026 Posit Software, PBC. All rights reserved. + * Licensed under the Elastic License 2.0. See LICENSE.txt for license information. + *--------------------------------------------------------------------------------------------*/ + +import * as vscode from 'vscode'; +import { randomUUID } from 'crypto'; +import { DATABRICKS_AUTH_PROVIDER_ID } from './constants'; +import { AuthProvider } from './authProvider'; +import { DatabricksLoopbackServer } from './databricksAuthServer'; +import { + buildAuthorizeUrl, + exchangeCodeForTokens, + generatePkcePair, + generateState, + normalizeHost, + refreshTokens, + TokenSet, +} from './databricksOAuth'; +import { log } from './log'; + +const SECRET_ACCESS_TOKEN = 'databricks.access_token'; +const SECRET_REFRESH_TOKEN = 'databricks.refresh_token'; +const SECRET_TOKEN_EXPIRY = 'databricks.token_expiry'; +const SECRET_HOST = 'databricks.host'; + +/** Refresh the access token when within this window of expiry. */ +const REFRESH_BUFFER_MS = 5 * 60 * 1000; + +/** How long to wait for the browser redirect before giving up. */ +const SIGN_IN_TIMEOUT_MS = 5 * 60 * 1000; + +/** + * Databricks authentication provider. + * + * Two credential paths: + * 1. OAuth U2M (desktop only) -- authorization code + PKCE against the + * built-in `databricks-cli` public client, with a loopback server on + * the fixed redirect port 8020. Tokens are refreshed lazily. + * 2. Personal access tokens -- the base-class API key machinery, used on + * remote/web where the loopback redirect cannot reach the extension + * host, or whenever the user prefers a PAT. + */ +export class DatabricksAuthProvider extends AuthProvider { + + /** Single in-flight refresh so concurrent getSessions calls don't double-refresh. */ + private _refreshPromise: Promise | null = null; + + private _signInCancellation: vscode.CancellationTokenSource | null = null; + + constructor(context: vscode.ExtensionContext) { + super(DATABRICKS_AUTH_PROVIDER_ID, 'Databricks', context); + } + + // --- AuthProvider overrides --- + + override async getSessions( + scopes?: readonly string[], + options?: vscode.AuthenticationProviderSessionOptions + ): Promise { + const sessions: vscode.AuthenticationSession[] = []; + + const oauthSession = await this.getOAuthSession(); + if (oauthSession && + (!options?.account || options.account.id === oauthSession.account.id)) { + sessions.push(oauthSession); + } + + // Stored personal access tokens (base-class machinery). + const patSessions = await super.getSessions(scopes, options); + return [...sessions, ...patSessions]; + } + + override async createSession( + _scopes: readonly string[], + _options?: vscode.AuthenticationProviderSessionOptions + ): Promise { + const host = normalizeHost(await this.resolveHost()); + await this.persistHostSetting(host); + + if (vscode.env.remoteName !== undefined || + vscode.env.uiKind === vscode.UIKind.Web) { + // The loopback redirect can't reach a remote extension host; + // fall back to a personal access token. + return this.createPatSession(host); + } + + return this.signInWithOAuth(host); + } + + override async removeSession(sessionId: string): Promise { + if (sessionId === DATABRICKS_AUTH_PROVIDER_ID) { + const removed = await this.buildStoredOAuthSession(); + await this.clearOAuthSecrets(); + if (removed) { + this.fireSessionsChanged({ + added: [], removed: [removed], changed: [], + }); + } + log.info('[Databricks] Signed out of OAuth session.'); + return; + } + return super.removeSession(sessionId); + } + + override cancelSignIn(): void { + this._signInCancellation?.cancel(); + this._signInCancellation?.dispose(); + this._signInCancellation = null; + } + + override dispose(): void { + this.cancelSignIn(); + super.dispose(); + } + + // --- OAuth session management --- + + /** + * Return the stored OAuth session, lazily refreshing the access token + * when it is within REFRESH_BUFFER_MS of expiry. Returns undefined when + * no OAuth credentials are stored or the refresh fails (the stored + * credentials are cleared in that case). + */ + private async getOAuthSession( + ): Promise { + const accessToken = await this.context.secrets.get(SECRET_ACCESS_TOKEN); + const expiry = await this.context.secrets.get(SECRET_TOKEN_EXPIRY); + if (!accessToken || !expiry) { + return undefined; + } + + let token: string | undefined = accessToken; + if (Date.now() >= parseInt(expiry) - REFRESH_BUFFER_MS) { + if (!this._refreshPromise) { + this._refreshPromise = this.refreshOAuthTokens() + .finally(() => { this._refreshPromise = null; }); + } + token = await this._refreshPromise; + if (!token) { + return undefined; + } + } + + const host = await this.context.secrets.get(SECRET_HOST); + return this.makeOAuthSession(token, host); + } + + /** + * Refresh the stored tokens. On failure, clears the OAuth secrets and + * fires a removed event so consumers know to re-authenticate. + */ + private async refreshOAuthTokens(): Promise { + try { + const refreshToken = await this.context.secrets.get(SECRET_REFRESH_TOKEN); + const host = await this.context.secrets.get(SECRET_HOST); + if (!refreshToken || !host) { + throw new Error('No stored refresh token or workspace host'); + } + log.info('[Databricks] Refreshing OAuth access token.'); + const tokens = await refreshTokens(host, refreshToken); + await this.storeOAuthSecrets(host, tokens); + log.info('[Databricks] OAuth access token refreshed.'); + return tokens.accessToken; + } catch (err) { + log.error(`[Databricks] Failed to refresh OAuth access token: ${err instanceof Error ? err.message : String(err)}`); + const removed = await this.buildStoredOAuthSession(); + await this.clearOAuthSecrets(); + if (removed) { + this.fireSessionsChanged({ + added: [], removed: [removed], changed: [], + }); + } + return undefined; + } + } + + private async signInWithOAuth( + host: string + ): Promise { + const state = generateState(); + const { verifier, challenge } = generatePkcePair(); + const server = new DatabricksLoopbackServer(state); + const cancellation = new vscode.CancellationTokenSource(); + this._signInCancellation = cancellation; + + try { + await server.start(); + const authorizeUrl = buildAuthorizeUrl(host, state, challenge); + log.info(`[Databricks] Starting OAuth sign-in for ${host}.`); + await vscode.env.openExternal(vscode.Uri.parse(authorizeUrl)); + + const code = await server.waitForCode( + SIGN_IN_TIMEOUT_MS, cancellation.token + ); + const tokens = await exchangeCodeForTokens(host, code, verifier); + await this.storeOAuthSecrets(host, tokens); + + const session = this.makeOAuthSession(tokens.accessToken, host); + this.fireSessionsChanged({ + added: [session], removed: [], changed: [], + }); + log.info('[Databricks] OAuth sign-in successful.'); + return session; + } finally { + await server.stop(); + cancellation.dispose(); + if (this._signInCancellation === cancellation) { + this._signInCancellation = null; + } + } + } + + private async createPatSession( + host: string + ): Promise { + const raw = await vscode.window.showInputBox({ + prompt: vscode.l10n.t( + 'Enter a Databricks personal access token for {0}', host + ), + password: true, + ignoreFocusOut: true, + validateInput: value => value.trim() + ? undefined + : vscode.l10n.t('A personal access token is required'), + }); + const token = raw?.trim(); + if (!token) { + throw new Error(vscode.l10n.t('A personal access token is required')); + } + return this.storeKey( + randomUUID(), this.accountLabel(host), token + ); + } + + // --- Helpers --- + + /** + * Resolve the workspace host: setting, then environment variable, + * then prompt the user. + */ + private async resolveHost(): Promise { + const credentials = vscode.workspace + .getConfiguration('authentication.databricks') + .get>('credentials', {}); + const configuredHost = credentials?.DATABRICKS_HOST?.trim(); + if (configuredHost) { + return configuredHost; + } + + const envHost = process.env.DATABRICKS_HOST?.trim(); + if (envHost) { + return envHost; + } + + const input = await vscode.window.showInputBox({ + prompt: vscode.l10n.t( + 'Enter your Databricks workspace URL (e.g. https://adb-1234567890123456.7.azuredatabricks.net)' + ), + ignoreFocusOut: true, + validateInput: value => value.trim() + ? undefined + : vscode.l10n.t('A workspace URL is required'), + }); + const host = input?.trim(); + if (!host) { + throw new Error(vscode.l10n.t('A Databricks workspace URL is required')); + } + return host; + } + + /** + * Persist the normalized host back to the global credentials setting. + * Read the global scope only so workspace-scoped values are not copied + * into global (same pattern as the Snowflake account sync). + */ + private async persistHostSetting(host: string): Promise { + const cfg = vscode.workspace + .getConfiguration('authentication.databricks'); + const inspection = cfg.inspect>('credentials'); + const globalValue = inspection?.globalValue ?? {}; + if (globalValue.DATABRICKS_HOST !== host) { + await cfg.update( + 'credentials', + { ...globalValue, DATABRICKS_HOST: host }, + vscode.ConfigurationTarget.Global + ).then(undefined, err => + log.error(`[Databricks] Failed to persist workspace host: ${err}`) + ); + } + } + + private accountLabel(host: string | undefined): string { + let hostname = host ?? ''; + try { + hostname = host ? new URL(normalizeHost(host)).hostname : ''; + } catch { + // Fall back to the raw host string. + } + return hostname ? `Databricks (${hostname})` : 'Databricks'; + } + + private makeOAuthSession( + accessToken: string, + host: string | undefined + ): vscode.AuthenticationSession { + return { + id: DATABRICKS_AUTH_PROVIDER_ID, + accessToken, + account: { + id: DATABRICKS_AUTH_PROVIDER_ID, + label: this.accountLabel(host), + }, + scopes: [], + }; + } + + /** + * Build a session object from the stored secrets, for removed events. + */ + private async buildStoredOAuthSession( + ): Promise { + const accessToken = await this.context.secrets.get(SECRET_ACCESS_TOKEN); + if (!accessToken) { + return undefined; + } + const host = await this.context.secrets.get(SECRET_HOST); + return this.makeOAuthSession(accessToken, host); + } + + private async storeOAuthSecrets( + host: string, + tokens: TokenSet + ): Promise { + await this.context.secrets.store(SECRET_ACCESS_TOKEN, tokens.accessToken); + await this.context.secrets.store(SECRET_REFRESH_TOKEN, tokens.refreshToken); + await this.context.secrets.store(SECRET_TOKEN_EXPIRY, tokens.expiresAt.toString()); + await this.context.secrets.store(SECRET_HOST, host); + } + + private async clearOAuthSecrets(): Promise { + await this.context.secrets.delete(SECRET_ACCESS_TOKEN); + await this.context.secrets.delete(SECRET_REFRESH_TOKEN); + await this.context.secrets.delete(SECRET_TOKEN_EXPIRY); + await this.context.secrets.delete(SECRET_HOST); + } +} diff --git a/extensions/authentication/src/databricksAuthServer.ts b/extensions/authentication/src/databricksAuthServer.ts new file mode 100644 index 000000000000..07c39f2d7d1a --- /dev/null +++ b/extensions/authentication/src/databricksAuthServer.ts @@ -0,0 +1,189 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (C) 2026 Posit Software, PBC. All rights reserved. + * Licensed under the Elastic License 2.0. See LICENSE.txt for license information. + *--------------------------------------------------------------------------------------------*/ + +import * as http from 'http'; +import type { CancellationToken } from 'vscode'; +import { DATABRICKS_OAUTH_REDIRECT_PORT } from './constants'; + +const DEFAULT_TIMEOUT_MS = 5 * 60 * 1000; + +const SUCCESS_HTML = ` + +Databricks Sign In + +

You are signed in to Databricks. You can close this tab.

+ +`; + +function errorHtml(message: string): string { + return ` + +Databricks Sign In + +

Databricks sign-in failed

+

${message}

+ +`; +} + +/** + * Minimal loopback HTTP server for the Databricks OAuth U2M flow. + * + * Databricks' built-in `databricks-cli` public client only allows the fixed + * redirect URI http://localhost:8020, so unlike the GitHub loopback server + * this one must bind a specific port and serves a tiny inline response + * instead of static media. The port is injectable for tests. + */ +export class DatabricksLoopbackServer { + private _server: http.Server | undefined; + private readonly _codePromise: Promise; + private _resolveCode!: (code: string) => void; + private _rejectCode!: (reason: Error) => void; + private _stopped = false; + + constructor( + private readonly expectedState: string, + private readonly port: number = DATABRICKS_OAUTH_REDIRECT_PORT, + ) { + this._codePromise = new Promise((resolve, reject) => { + this._resolveCode = resolve; + this._rejectCode = reject; + }); + // Avoid an unhandled rejection if the promise settles before + // waitForCode attaches handlers (e.g. an early bad request). + this._codePromise.catch(() => { }); + } + + /** + * Start listening on 127.0.0.1:. + */ + start(): Promise { + return new Promise((resolve, reject) => { + if (this._server) { + reject(new Error('Server is already started')); + return; + } + const server = http.createServer( + (req, res) => this.handleRequest(req, res) + ); + this._server = server; + server.on('error', (err: NodeJS.ErrnoException) => { + if (err.code === 'EADDRINUSE') { + reject(new Error( + `Port ${this.port} is already in use. Databricks ` + + `sign-in requires port ${this.port} - close the ` + + 'application using it, or use a personal access ' + + 'token instead.' + )); + } else { + reject(new Error(`Failed to start sign-in server: ${err.message}`)); + } + }); + server.listen(this.port, '127.0.0.1', () => resolve()); + }); + } + + private handleRequest( + req: http.IncomingMessage, + res: http.ServerResponse + ): void { + // Accept GET on any path; Databricks redirects to the URI root. + const reqUrl = new URL(req.url ?? '/', `http://127.0.0.1:${this.port}`); + const code = reqUrl.searchParams.get('code'); + const state = reqUrl.searchParams.get('state'); + const error = reqUrl.searchParams.get('error'); + const errorDescription = reqUrl.searchParams.get('error_description'); + + // Browsers request /favicon.ico alongside the redirect; ignore + // requests that carry no OAuth response parameters. + if (!code && !state && !error) { + res.writeHead(404); + res.end(); + return; + } + + if (error) { + const message = errorDescription ?? error; + res.writeHead(400, { 'Content-Type': 'text/html' }); + res.end(errorHtml(message)); + this._rejectCode(new Error(message)); + return; + } + + if (state !== this.expectedState) { + res.writeHead(400, { 'Content-Type': 'text/html' }); + res.end(errorHtml('State mismatch. Please try signing in again.')); + this._rejectCode(new Error( + 'Databricks sign-in failed: state parameter does not match.' + )); + return; + } + + if (!code) { + res.writeHead(400, { 'Content-Type': 'text/html' }); + res.end(errorHtml('Missing authorization code.')); + this._rejectCode(new Error( + 'Databricks sign-in failed: no authorization code received.' + )); + return; + } + + res.writeHead(200, { 'Content-Type': 'text/html' }); + res.end(SUCCESS_HTML); + this._resolveCode(code); + } + + /** + * Wait for the authorization code, racing a timeout and an optional + * cancellation token. + */ + waitForCode( + timeoutMs: number = DEFAULT_TIMEOUT_MS, + cancellationToken?: CancellationToken + ): Promise { + return new Promise((resolve, reject) => { + const timeout = setTimeout(() => { + reject(new Error( + 'Timed out waiting for Databricks sign-in. Please try again.' + )); + }, timeoutMs); + + const cancellation = cancellationToken?.onCancellationRequested(() => { + clearTimeout(timeout); + reject(new Error('Databricks sign-in was cancelled.')); + }); + + this._codePromise.then( + code => { + clearTimeout(timeout); + cancellation?.dispose(); + resolve(code); + }, + err => { + clearTimeout(timeout); + cancellation?.dispose(); + reject(err); + } + ); + }); + } + + /** + * Stop the server. Safe to call multiple times. + */ + stop(): Promise { + return new Promise((resolve) => { + if (this._stopped || !this._server) { + this._stopped = true; + resolve(); + return; + } + this._stopped = true; + this._server.close(() => resolve()); + // Close keep-alive connections so close() completes promptly. + this._server.closeAllConnections?.(); + }); + } +} diff --git a/extensions/authentication/src/databricksOAuth.ts b/extensions/authentication/src/databricksOAuth.ts new file mode 100644 index 000000000000..d74661305143 --- /dev/null +++ b/extensions/authentication/src/databricksOAuth.ts @@ -0,0 +1,161 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (C) 2026 Posit Software, PBC. All rights reserved. + * Licensed under the Elastic License 2.0. See LICENSE.txt for license information. + *--------------------------------------------------------------------------------------------*/ + +import { createHash, randomBytes } from 'crypto'; +import { + DATABRICKS_OAUTH_CLIENT_ID, + DATABRICKS_OAUTH_REDIRECT_URI, + DATABRICKS_OAUTH_SCOPES, +} from './constants'; + +/** + * Pure helpers for the Databricks OAuth 2.0 U2M flow (authorization code + + * PKCE against the built-in `databricks-cli` public client). No vscode + * imports so these can be unit tested directly. + */ + +/** A resolved set of OAuth tokens. `expiresAt` is epoch milliseconds. */ +export interface TokenSet { + accessToken: string; + refreshToken: string; + expiresAt: number; +} + +interface TokenEndpointResponse { + access_token: string; + refresh_token?: string; + expires_in: number; + token_type?: string; +} + +/** + * Generate a PKCE verifier/challenge pair (S256). + */ +export function generatePkcePair(): { verifier: string; challenge: string } { + const verifier = randomBytes(32).toString('base64url'); + const challenge = createHash('sha256').update(verifier).digest('base64url'); + return { verifier, challenge }; +} + +/** + * Generate an opaque state value for CSRF protection. + */ +export function generateState(): string { + return randomBytes(16).toString('base64url'); +} + +/** + * Normalize a user-supplied workspace host: trim whitespace, prepend + * https:// when no scheme is present, and strip trailing slashes. + */ +export function normalizeHost(raw: string): string { + let host = raw.trim(); + if (host && !/^[a-zA-Z][a-zA-Z0-9+.-]*:\/\//.test(host)) { + host = `https://${host}`; + } + return host.replace(/\/+$/, ''); +} + +/** + * Build the authorization URL for the Databricks OIDC authorize endpoint. + */ +export function buildAuthorizeUrl( + host: string, + state: string, + challenge: string +): string { + const url = new URL(`${normalizeHost(host)}/oidc/v1/authorize`); + url.searchParams.set('client_id', DATABRICKS_OAUTH_CLIENT_ID); + url.searchParams.set('response_type', 'code'); + url.searchParams.set('redirect_uri', DATABRICKS_OAUTH_REDIRECT_URI); + url.searchParams.set('scope', DATABRICKS_OAUTH_SCOPES); + url.searchParams.set('state', state); + url.searchParams.set('code_challenge', challenge); + url.searchParams.set('code_challenge_method', 'S256'); + return url.toString(); +} + +async function postTokenEndpoint( + host: string, + params: Record, + operation: string +): Promise { + const response = await fetch(`${normalizeHost(host)}/oidc/v1/token`, { + method: 'POST', + headers: { 'Content-Type': 'application/x-www-form-urlencoded' }, + body: new URLSearchParams(params).toString(), + }); + + if (!response.ok) { + const errorData = await response.json().catch(() => ({})) as { + error?: string; + error_description?: string; + }; + const detail = errorData.error_description + ?? errorData.error + ?? response.statusText; + throw new Error( + `Databricks ${operation} failed (HTTP ${response.status}): ${detail}` + ); + } + + return await response.json() as TokenEndpointResponse; +} + +/** + * Exchange an authorization code for tokens. + */ +export async function exchangeCodeForTokens( + host: string, + code: string, + verifier: string +): Promise { + const data = await postTokenEndpoint(host, { + grant_type: 'authorization_code', + code, + redirect_uri: DATABRICKS_OAUTH_REDIRECT_URI, + client_id: DATABRICKS_OAUTH_CLIENT_ID, + code_verifier: verifier, + }, 'token exchange'); + + if (!data.access_token || !data.refresh_token) { + throw new Error( + 'Databricks token exchange response is missing tokens' + ); + } + + return { + accessToken: data.access_token, + refreshToken: data.refresh_token, + expiresAt: Date.now() + data.expires_in * 1000, + }; +} + +/** + * Refresh an access token. Keeps the existing refresh token unless the + * response includes a rotated one. + */ +export async function refreshTokens( + host: string, + refreshToken: string +): Promise { + const data = await postTokenEndpoint(host, { + grant_type: 'refresh_token', + refresh_token: refreshToken, + client_id: DATABRICKS_OAUTH_CLIENT_ID, + }, 'token refresh'); + + if (!data.access_token) { + throw new Error( + 'Databricks token refresh response is missing an access token' + ); + } + + return { + accessToken: data.access_token, + refreshToken: data.refresh_token ?? refreshToken, + expiresAt: Date.now() + data.expires_in * 1000, + }; +} diff --git a/extensions/authentication/src/extension.ts b/extensions/authentication/src/extension.ts index c43ea047fe77..bfb1562007ca 100644 --- a/extensions/authentication/src/extension.ts +++ b/extensions/authentication/src/extension.ts @@ -11,6 +11,7 @@ import { AWS_AUTH_PROVIDER_ID, CREDENTIAL_REFRESH_INTERVAL_MS, CUSTOM_PROVIDER_AUTH_PROVIDER_ID, + DATABRICKS_AUTH_PROVIDER_ID, DEEPSEEK_AUTH_PROVIDER_ID, FOUNDRY_AUTH_PROVIDER_ID, GEMINI_AUTH_PROVIDER_ID, @@ -25,6 +26,7 @@ import { normalizeToV1Url, validateAnthropicApiKey, validateCustomProviderApiKey, + validateDatabricksApiKey, validateDeepSeekApiKey, validateFoundryApiKey, validateGeminiApiKey, @@ -39,6 +41,8 @@ import { getSnowflakeConnectionsTomlPath, } from './credentials/snowflake'; import { PositOAuthProvider } from './positOAuthProvider'; +import { DatabricksAuthProvider } from './databricksAuthProvider'; +import { normalizeHost } from './databricksOAuth'; import * as fs from 'fs'; import { log } from './log'; import { migrateAwsSettings } from './migration/aws'; @@ -67,6 +71,7 @@ export async function activate(context: vscode.ExtensionContext) { await registerGeminiProvider(context); await registerGeapProvider(context); await registerDeepSeekProvider(context); + registerDatabricksProvider(context); registerCustomProvider(context); // Register providers so the Settings UI shows per-provider @@ -683,6 +688,46 @@ async function registerDeepSeekProvider( log.info(`Registered auth provider: ${DEEPSEEK_AUTH_PROVIDER_ID}`); } +function registerDatabricksProvider( + context: vscode.ExtensionContext +): void { + const logger = new AuthProviderLogger('Databricks'); + const provider = new DatabricksAuthProvider(context); + context.subscriptions.push( + vscode.authentication.registerAuthenticationProvider( + DATABRICKS_AUTH_PROVIDER_ID, 'Databricks', provider, + { supportsMultipleAccounts: false } + ), + provider + ); + registerAuthProvider(DATABRICKS_AUTH_PROVIDER_ID, provider, { + validateApiKey: validateDatabricksApiKey, + onSave: async (config) => { + // baseUrl holds the workspace host; persist it as + // DATABRICKS_HOST in the credentials setting. Read the global + // scope only so workspace-scoped values are not copied into + // global (same pattern as the Snowflake account sync). + const host = config.baseUrl?.trim(); + if (!host) { + return; + } + const normalized = normalizeHost(host); + const cfg = vscode.workspace + .getConfiguration('authentication.databricks'); + const inspection = cfg.inspect>('credentials'); + const globalValue = inspection?.globalValue ?? {}; + if (globalValue.DATABRICKS_HOST !== normalized) { + await cfg.update( + 'credentials', + { ...globalValue, DATABRICKS_HOST: normalized }, + vscode.ConfigurationTarget.Global + ); + } + }, + }); + logger.info('Registered auth provider'); +} + function registerCustomProvider( context: vscode.ExtensionContext ): void { diff --git a/extensions/authentication/src/positOAuthProvider.ts b/extensions/authentication/src/positOAuthProvider.ts index 3a994a910722..7c59bce049df 100644 --- a/extensions/authentication/src/positOAuthProvider.ts +++ b/extensions/authentication/src/positOAuthProvider.ts @@ -140,7 +140,7 @@ export class PositOAuthProvider extends AuthProvider { await this.context.secrets.delete('posit-ai.token_expiry'); } - cancelSignIn(): void { + override cancelSignIn(): void { this._cancellationToken?.cancel(); this._cancellationToken?.dispose(); this._cancellationToken = null; diff --git a/extensions/authentication/src/providerSources.ts b/extensions/authentication/src/providerSources.ts index a6173896f5c4..996f38190ea1 100644 --- a/extensions/authentication/src/providerSources.ts +++ b/extensions/authentication/src/providerSources.ts @@ -9,6 +9,7 @@ import { ANTHROPIC_AUTH_PROVIDER_ID, AWS_AUTH_PROVIDER_ID, CUSTOM_PROVIDER_AUTH_PROVIDER_ID, + DATABRICKS_AUTH_PROVIDER_ID, DEEPSEEK_AUTH_PROVIDER_ID, FOUNDRY_AUTH_PROVIDER_ID, GEMINI_AUTH_PROVIDER_ID, @@ -99,6 +100,11 @@ export const PROVIDER_METADATA: Record = { settingName: 'deepseek', status: 'experimental', }, + databricks: { + id: DATABRICKS_AUTH_PROVIDER_ID, + displayName: 'Databricks', + settingName: 'databricks', + }, }; export function getProviderSources(): positron.ai.LanguageModelSource[] { @@ -108,6 +114,14 @@ export function getProviderSources(): positron.ai.LanguageModelSource[] { const geapFromEnv = !!process.env.GOOGLE_VERTEX_PROJECT && !!process.env.GOOGLE_VERTEX_LOCATION; + // Databricks OAuth needs a loopback server on the fixed port 8020, + // which only works on desktop. Remote/web sessions use a PAT instead. + const databricksOauthAvailable = vscode.env.remoteName === undefined + && vscode.env.uiKind !== vscode.UIKind.Web; + const databricksHost = vscode.workspace + .getConfiguration('authentication.databricks') + .get>('credentials', {})?.DATABRICKS_HOST ?? ''; + return [ { type: positron.PositronLanguageModelType.Chat, @@ -252,5 +266,18 @@ export function getProviderSources(): positron.ai.LanguageModelSource[] { }, }, }, + { + type: positron.PositronLanguageModelType.Chat, + provider: PROVIDER_METADATA.databricks, + // baseUrl holds the Databricks workspace URL. + supportedOptions: databricksOauthAvailable + ? ['oauth', 'apiKey', 'baseUrl'] + : ['apiKey', 'baseUrl'], + defaults: { + model: 'databricks', + baseUrl: databricksHost, + toolCalls: true, + }, + }, ]; } diff --git a/extensions/authentication/src/test/databricksAuthProvider.test.ts b/extensions/authentication/src/test/databricksAuthProvider.test.ts new file mode 100644 index 000000000000..ca391bcbd76c --- /dev/null +++ b/extensions/authentication/src/test/databricksAuthProvider.test.ts @@ -0,0 +1,212 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (C) 2026 Posit Software, PBC. All rights reserved. + * Licensed under the Elastic License 2.0. See LICENSE.txt for license information. + *--------------------------------------------------------------------------------------------*/ + +import * as assert from 'assert'; +import * as vscode from 'vscode'; +import { DatabricksAuthProvider } from '../databricksAuthProvider'; + +const HOST = 'https://example.cloud.databricks.com'; + +function storeOAuthSecrets(secrets: Map, overrides?: { + accessToken?: string; + refreshToken?: string; + expiresAt?: number; + host?: string; +}): void { + secrets.set('databricks.access_token', overrides?.accessToken ?? 'test-access-token'); + secrets.set('databricks.refresh_token', overrides?.refreshToken ?? 'test-refresh-token'); + secrets.set('databricks.token_expiry', String(overrides?.expiresAt ?? Date.now() + 3600 * 1000)); + secrets.set('databricks.host', overrides?.host ?? HOST); +} + +function makeMockContext(): { context: vscode.ExtensionContext; secrets: Map } { + const secrets = new Map(); + const globalState = new Map(); + const context = { + secrets: { + get: (key: string) => Promise.resolve(secrets.get(key)), + store: (key: string, value: string) => { + secrets.set(key, value); + return Promise.resolve(); + }, + delete: (key: string) => { + secrets.delete(key); + return Promise.resolve(); + }, + }, + globalState: { + get: (key: string) => globalState.get(key) as T | undefined, + update: (key: string, value: unknown) => { + globalState.set(key, value); + return Promise.resolve(); + }, + }, + } as unknown as vscode.ExtensionContext; + return { context, secrets }; +} + +suite('DatabricksAuthProvider', () => { + let provider: DatabricksAuthProvider; + let secrets: Map; + let originalFetch: typeof globalThis.fetch; + + setup(() => { + originalFetch = globalThis.fetch; + const mock = makeMockContext(); + secrets = mock.secrets; + provider = new DatabricksAuthProvider(mock.context); + }); + + teardown(() => { + provider.dispose(); + globalThis.fetch = originalFetch; + }); + + suite('getSessions', () => { + test('returns the OAuth session when the token is fresh', async () => { + storeOAuthSecrets(secrets, { + accessToken: 'fresh-token', + expiresAt: Date.now() + 30 * 60 * 1000, + }); + + const sessions = await provider.getSessions(); + assert.strictEqual(sessions.length, 1); + assert.strictEqual(sessions[0].id, 'databricks'); + assert.strictEqual(sessions[0].accessToken, 'fresh-token'); + assert.strictEqual( + sessions[0].account.label, + 'Databricks (example.cloud.databricks.com)' + ); + }); + + test('returns no sessions when nothing is stored', async () => { + const sessions = await provider.getSessions(); + assert.deepStrictEqual(sessions, []); + }); + + test('refreshes a stale token and stores the result', async () => { + storeOAuthSecrets(secrets, { + accessToken: 'stale-token', + expiresAt: Date.now() + 60 * 1000, // within the 5 min buffer + }); + + globalThis.fetch = async () => new Response(JSON.stringify({ + access_token: 'new-token', + refresh_token: 'new-refresh', + expires_in: 3600, + }), { status: 200 }); + + const sessions = await provider.getSessions(); + assert.strictEqual(sessions.length, 1); + assert.strictEqual(sessions[0].accessToken, 'new-token'); + assert.strictEqual(secrets.get('databricks.access_token'), 'new-token'); + assert.strictEqual(secrets.get('databricks.refresh_token'), 'new-refresh'); + assert.strictEqual(secrets.get('databricks.host'), HOST); + }); + + test('clears secrets and fires removed when refresh fails', async () => { + storeOAuthSecrets(secrets, { expiresAt: Date.now() - 1000 }); + + globalThis.fetch = async () => new Response(JSON.stringify({ + error: 'invalid_grant', + error_description: 'Refresh token revoked', + }), { status: 401 }); + + const events: vscode.AuthenticationProviderAuthenticationSessionsChangeEvent[] = []; + const subscription = provider.onDidChangeSessions(e => events.push(e)); + try { + const sessions = await provider.getSessions(); + assert.deepStrictEqual(sessions, []); + } finally { + subscription.dispose(); + } + + assert.strictEqual(secrets.get('databricks.access_token'), undefined); + assert.strictEqual(secrets.get('databricks.refresh_token'), undefined); + assert.strictEqual(secrets.get('databricks.token_expiry'), undefined); + assert.strictEqual(secrets.get('databricks.host'), undefined); + + assert.strictEqual(events.length, 1); + assert.strictEqual(events[0].removed?.length, 1); + assert.strictEqual(events[0].removed?.[0].id, 'databricks'); + }); + + test('concurrent calls share a single refresh', async () => { + storeOAuthSecrets(secrets, { expiresAt: Date.now() - 1000 }); + + let fetchCount = 0; + globalThis.fetch = async () => { + fetchCount++; + // Yield so both getSessions calls overlap the refresh. + await new Promise(resolve => setTimeout(resolve, 20)); + return new Response(JSON.stringify({ + access_token: 'new-token', + refresh_token: 'new-refresh', + expires_in: 3600, + }), { status: 200 }); + }; + + const [a, b] = await Promise.all([ + provider.getSessions(), + provider.getSessions(), + ]); + + assert.strictEqual(fetchCount, 1); + assert.strictEqual(a[0].accessToken, 'new-token'); + assert.strictEqual(b[0].accessToken, 'new-token'); + }); + + test('falls through to stored PAT sessions', async () => { + await provider.storeKey('pat-account', 'Databricks (PAT)', 'dapi123'); + + const sessions = await provider.getSessions(); + assert.strictEqual(sessions.length, 1); + assert.strictEqual(sessions[0].id, 'pat-account'); + assert.strictEqual(sessions[0].accessToken, 'dapi123'); + }); + + test('returns both OAuth and PAT sessions', async () => { + storeOAuthSecrets(secrets, { + expiresAt: Date.now() + 30 * 60 * 1000, + }); + await provider.storeKey('pat-account', 'Databricks (PAT)', 'dapi123'); + + const sessions = await provider.getSessions(); + assert.strictEqual(sessions.length, 2); + assert.strictEqual(sessions[0].id, 'databricks'); + assert.strictEqual(sessions[1].id, 'pat-account'); + }); + }); + + suite('removeSession', () => { + test('clears OAuth secrets and fires removed for the OAuth session', async () => { + storeOAuthSecrets(secrets); + + const events: vscode.AuthenticationProviderAuthenticationSessionsChangeEvent[] = []; + const subscription = provider.onDidChangeSessions(e => events.push(e)); + try { + await provider.removeSession('databricks'); + } finally { + subscription.dispose(); + } + + assert.strictEqual(secrets.get('databricks.access_token'), undefined); + assert.strictEqual(secrets.get('databricks.refresh_token'), undefined); + assert.strictEqual(secrets.get('databricks.token_expiry'), undefined); + assert.strictEqual(secrets.get('databricks.host'), undefined); + assert.strictEqual(events.length, 1); + assert.strictEqual(events[0].removed?.[0].id, 'databricks'); + }); + + test('delegates PAT sessions to the base class', async () => { + await provider.storeKey('pat-account', 'Databricks (PAT)', 'dapi123'); + + await provider.removeSession('pat-account'); + + const sessions = await provider.getSessions(); + assert.deepStrictEqual(sessions, []); + }); + }); +}); diff --git a/extensions/authentication/src/test/databricksAuthServer.test.ts b/extensions/authentication/src/test/databricksAuthServer.test.ts new file mode 100644 index 000000000000..155b91536f8a --- /dev/null +++ b/extensions/authentication/src/test/databricksAuthServer.test.ts @@ -0,0 +1,139 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (C) 2026 Posit Software, PBC. All rights reserved. + * Licensed under the Elastic License 2.0. See LICENSE.txt for license information. + *--------------------------------------------------------------------------------------------*/ + +import * as assert from 'assert'; +import * as http from 'http'; +import * as net from 'net'; +import { DatabricksLoopbackServer } from '../databricksAuthServer'; + +/** Find a free port by briefly binding an ephemeral listener. */ +function getFreePort(): Promise { + return new Promise((resolve, reject) => { + const probe = net.createServer(); + probe.once('error', reject); + probe.listen(0, '127.0.0.1', () => { + const port = (probe.address() as net.AddressInfo).port; + probe.close(() => resolve(port)); + }); + }); +} + +function get(port: number, pathAndQuery: string): Promise<{ status: number; body: string }> { + return new Promise((resolve, reject) => { + http.get(`http://127.0.0.1:${port}${pathAndQuery}`, res => { + let body = ''; + res.on('data', chunk => { body += chunk; }); + res.on('end', () => resolve({ status: res.statusCode ?? 0, body })); + }).on('error', reject); + }); +} + +suite('DatabricksLoopbackServer', () => { + let server: DatabricksLoopbackServer | undefined; + + teardown(async () => { + await server?.stop(); + server = undefined; + }); + + test('resolves the code on a valid redirect', async () => { + const port = await getFreePort(); + server = new DatabricksLoopbackServer('expected-state', port); + await server.start(); + + const codePromise = server.waitForCode(5000); + const response = await get(port, '/?code=auth-code-123&state=expected-state'); + + assert.strictEqual(response.status, 200); + assert.ok(response.body.includes('You are signed in to Databricks')); + assert.strictEqual(await codePromise, 'auth-code-123'); + }); + + test('accepts the redirect on any path', async () => { + const port = await getFreePort(); + server = new DatabricksLoopbackServer('expected-state', port); + await server.start(); + + const codePromise = server.waitForCode(5000); + const response = await get(port, '/some/path?code=abc&state=expected-state'); + + assert.strictEqual(response.status, 200); + assert.strictEqual(await codePromise, 'abc'); + }); + + test('responds 400 and rejects on a state mismatch', async () => { + const port = await getFreePort(); + server = new DatabricksLoopbackServer('expected-state', port); + await server.start(); + + const codePromise = server.waitForCode(5000); + const response = await get(port, '/?code=auth-code-123&state=wrong-state'); + + assert.strictEqual(response.status, 400); + await assert.rejects( + () => codePromise, + (err: Error) => err.message.includes('state') + ); + }); + + test('rejects with the error description on an error redirect', async () => { + const port = await getFreePort(); + server = new DatabricksLoopbackServer('expected-state', port); + await server.start(); + + const codePromise = server.waitForCode(5000); + const response = await get( + port, + '/?error=access_denied&error_description=User%20denied%20access&state=expected-state' + ); + + assert.strictEqual(response.status, 400); + await assert.rejects( + () => codePromise, + (err: Error) => err.message === 'User denied access' + ); + }); + + test('start maps EADDRINUSE to a friendly error', async () => { + const port = await getFreePort(); + // Pre-bind the port so start() fails. + const blocker = net.createServer(); + await new Promise((resolve, reject) => { + blocker.once('error', reject); + blocker.listen(port, '127.0.0.1', () => resolve()); + }); + + try { + server = new DatabricksLoopbackServer('expected-state', port); + await assert.rejects( + () => server!.start(), + (err: Error) => + err.message.includes(`Port ${port} is already in use`) && + err.message.includes('personal access token') + ); + } finally { + await new Promise(resolve => blocker.close(() => resolve())); + } + }); + + test('waitForCode times out', async () => { + const port = await getFreePort(); + server = new DatabricksLoopbackServer('expected-state', port); + await server.start(); + + await assert.rejects( + () => server!.waitForCode(50), + (err: Error) => err.message.includes('Timed out') + ); + }); + + test('stop is idempotent', async () => { + const port = await getFreePort(); + server = new DatabricksLoopbackServer('expected-state', port); + await server.start(); + await server.stop(); + await server.stop(); + }); +}); diff --git a/extensions/authentication/src/test/databricksOAuth.test.ts b/extensions/authentication/src/test/databricksOAuth.test.ts new file mode 100644 index 000000000000..4df81b4a795a --- /dev/null +++ b/extensions/authentication/src/test/databricksOAuth.test.ts @@ -0,0 +1,246 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (C) 2026 Posit Software, PBC. All rights reserved. + * Licensed under the Elastic License 2.0. See LICENSE.txt for license information. + *--------------------------------------------------------------------------------------------*/ + +import * as assert from 'assert'; +import { createHash } from 'crypto'; +import { + buildAuthorizeUrl, + exchangeCodeForTokens, + generatePkcePair, + generateState, + normalizeHost, + refreshTokens, +} from '../databricksOAuth'; + +const BASE64URL_PATTERN = /^[A-Za-z0-9_-]+$/; + +function mockFetch( + handler: (url: string, init?: RequestInit) => Response +): { calls: { url: string; init?: RequestInit }[] } { + const calls: { url: string; init?: RequestInit }[] = []; + globalThis.fetch = async (input: string | URL | Request, init?: RequestInit) => { + const url = String(input); + calls.push({ url, init }); + return handler(url, init); + }; + return { calls }; +} + +function tokenResponse(body: unknown, status = 200): Response { + return new Response(JSON.stringify(body), { status }); +} + +suite('databricksOAuth', () => { + let originalFetch: typeof globalThis.fetch; + + setup(() => { + originalFetch = globalThis.fetch; + }); + + teardown(() => { + globalThis.fetch = originalFetch; + }); + + suite('generatePkcePair', () => { + test('verifier is 32 bytes of base64url', () => { + const { verifier } = generatePkcePair(); + // 32 bytes -> 43 base64url chars, no padding. + assert.strictEqual(verifier.length, 43); + assert.match(verifier, BASE64URL_PATTERN); + }); + + test('challenge is base64url(sha256(verifier))', () => { + const { verifier, challenge } = generatePkcePair(); + const expected = createHash('sha256').update(verifier).digest('base64url'); + assert.strictEqual(challenge, expected); + assert.match(challenge, BASE64URL_PATTERN); + }); + + test('pairs are unique', () => { + const a = generatePkcePair(); + const b = generatePkcePair(); + assert.notStrictEqual(a.verifier, b.verifier); + }); + }); + + suite('generateState', () => { + test('is 16 bytes of base64url', () => { + const state = generateState(); + // 16 bytes -> 22 base64url chars, no padding. + assert.strictEqual(state.length, 22); + assert.match(state, BASE64URL_PATTERN); + }); + }); + + suite('normalizeHost', () => { + test('trims whitespace', () => { + assert.strictEqual( + normalizeHost(' https://example.cloud.databricks.com '), + 'https://example.cloud.databricks.com' + ); + }); + + test('prepends https:// when no scheme', () => { + assert.strictEqual( + normalizeHost('example.cloud.databricks.com'), + 'https://example.cloud.databricks.com' + ); + }); + + test('keeps an existing scheme', () => { + assert.strictEqual( + normalizeHost('http://localhost:8080'), + 'http://localhost:8080' + ); + }); + + test('strips trailing slashes', () => { + assert.strictEqual( + normalizeHost('https://example.cloud.databricks.com///'), + 'https://example.cloud.databricks.com' + ); + }); + + test('combines all normalizations', () => { + assert.strictEqual( + normalizeHost(' example.cloud.databricks.com/ '), + 'https://example.cloud.databricks.com' + ); + }); + }); + + suite('buildAuthorizeUrl', () => { + test('builds the authorize URL with exact params', () => { + const url = new URL(buildAuthorizeUrl( + 'https://example.cloud.databricks.com', 'the-state', 'the-challenge' + )); + assert.strictEqual(url.origin, 'https://example.cloud.databricks.com'); + assert.strictEqual(url.pathname, '/oidc/v1/authorize'); + assert.strictEqual(url.searchParams.get('client_id'), 'databricks-cli'); + assert.strictEqual(url.searchParams.get('response_type'), 'code'); + assert.strictEqual(url.searchParams.get('redirect_uri'), 'http://localhost:8020'); + assert.strictEqual(url.searchParams.get('scope'), 'all-apis offline_access'); + assert.strictEqual(url.searchParams.get('state'), 'the-state'); + assert.strictEqual(url.searchParams.get('code_challenge'), 'the-challenge'); + assert.strictEqual(url.searchParams.get('code_challenge_method'), 'S256'); + assert.strictEqual([...url.searchParams.keys()].length, 7); + }); + + test('normalizes the host', () => { + const url = buildAuthorizeUrl('example.cloud.databricks.com/', 's', 'c'); + assert.ok(url.startsWith('https://example.cloud.databricks.com/oidc/v1/authorize?')); + }); + }); + + suite('exchangeCodeForTokens', () => { + test('posts the form-encoded exchange request and computes expiry', async () => { + const before = Date.now(); + const { calls } = mockFetch(() => tokenResponse({ + access_token: 'access-1', + refresh_token: 'refresh-1', + expires_in: 3600, + token_type: 'Bearer', + })); + + const tokens = await exchangeCodeForTokens( + 'https://example.cloud.databricks.com', 'the-code', 'the-verifier' + ); + + assert.strictEqual(calls.length, 1); + assert.strictEqual(calls[0].url, 'https://example.cloud.databricks.com/oidc/v1/token'); + assert.strictEqual(calls[0].init?.method, 'POST'); + assert.strictEqual( + (calls[0].init?.headers as Record)['Content-Type'], + 'application/x-www-form-urlencoded' + ); + const body = new URLSearchParams(calls[0].init?.body as string); + assert.strictEqual(body.get('grant_type'), 'authorization_code'); + assert.strictEqual(body.get('code'), 'the-code'); + assert.strictEqual(body.get('redirect_uri'), 'http://localhost:8020'); + assert.strictEqual(body.get('client_id'), 'databricks-cli'); + assert.strictEqual(body.get('code_verifier'), 'the-verifier'); + + assert.strictEqual(tokens.accessToken, 'access-1'); + assert.strictEqual(tokens.refreshToken, 'refresh-1'); + assert.ok(tokens.expiresAt >= before + 3600 * 1000); + assert.ok(tokens.expiresAt <= Date.now() + 3600 * 1000); + }); + + test('throws an informative error on non-200', async () => { + mockFetch(() => tokenResponse({ + error: 'invalid_grant', + error_description: 'Authorization code expired', + }, 400)); + + await assert.rejects( + () => exchangeCodeForTokens('https://example.com', 'code', 'verifier'), + (err: Error) => + err.message.includes('token exchange') && + err.message.includes('400') && + err.message.includes('Authorization code expired') + ); + }); + + test('throws on non-200 with a non-JSON body', async () => { + mockFetch(() => new Response('Bad Gateway', { + status: 502, statusText: 'Bad Gateway', + })); + + await assert.rejects( + () => exchangeCodeForTokens('https://example.com', 'code', 'verifier'), + (err: Error) => err.message.includes('502') + ); + }); + }); + + suite('refreshTokens', () => { + test('posts the form-encoded refresh request', async () => { + const { calls } = mockFetch(() => tokenResponse({ + access_token: 'access-2', + refresh_token: 'refresh-2', + expires_in: 1800, + })); + + const tokens = await refreshTokens( + 'https://example.cloud.databricks.com', 'refresh-1' + ); + + assert.strictEqual(calls[0].url, 'https://example.cloud.databricks.com/oidc/v1/token'); + const body = new URLSearchParams(calls[0].init?.body as string); + assert.strictEqual(body.get('grant_type'), 'refresh_token'); + assert.strictEqual(body.get('refresh_token'), 'refresh-1'); + assert.strictEqual(body.get('client_id'), 'databricks-cli'); + assert.strictEqual([...body.keys()].length, 3); + + // Rotated refresh token is returned. + assert.strictEqual(tokens.accessToken, 'access-2'); + assert.strictEqual(tokens.refreshToken, 'refresh-2'); + }); + + test('keeps the old refresh token when the response omits one', async () => { + mockFetch(() => tokenResponse({ + access_token: 'access-2', + expires_in: 1800, + })); + + const tokens = await refreshTokens('https://example.com', 'refresh-1'); + assert.strictEqual(tokens.refreshToken, 'refresh-1'); + }); + + test('throws an informative error on non-200', async () => { + mockFetch(() => tokenResponse({ + error: 'invalid_grant', + error_description: 'Refresh token revoked', + }, 401)); + + await assert.rejects( + () => refreshTokens('https://example.com', 'refresh-1'), + (err: Error) => + err.message.includes('token refresh') && + err.message.includes('Refresh token revoked') + ); + }); + }); +}); diff --git a/extensions/authentication/src/validation/databricks.ts b/extensions/authentication/src/validation/databricks.ts new file mode 100644 index 000000000000..544b67efef47 --- /dev/null +++ b/extensions/authentication/src/validation/databricks.ts @@ -0,0 +1,94 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (C) 2026 Posit Software, PBC. All rights reserved. + * Licensed under the Elastic License 2.0. See LICENSE.txt for license information. + *--------------------------------------------------------------------------------------------*/ + +import * as vscode from 'vscode'; +import * as positron from 'positron'; +import { KEY_VALIDATION_TIMEOUT_MS } from '../constants'; +import { normalizeHost } from '../databricksOAuth'; + +class DatabricksValidationError extends Error { + constructor(message: string) { + super(message); + this.name = 'DatabricksValidationError'; + } +} + +/** + * Resolve the workspace host for validation: the config dialog's baseUrl + * field, falling back to the saved credentials setting. + */ +function resolveValidationHost( + config: positron.ai.LanguageModelConfig +): string | undefined { + const fromConfig = config.baseUrl?.trim(); + if (fromConfig) { + return fromConfig; + } + const credentials = vscode.workspace + .getConfiguration('authentication.databricks') + .get>('credentials', {}); + return credentials?.DATABRICKS_HOST?.trim() || undefined; +} + +/** + * Validate a Databricks personal access token against the workspace's + * SCIM Me endpoint. + */ +export async function validateDatabricksApiKey( + apiKey: string, + config: positron.ai.LanguageModelConfig +): Promise { + const rawHost = resolveValidationHost(config); + if (!rawHost) { + throw new DatabricksValidationError(vscode.l10n.t( + 'Databricks workspace URL is required (e.g. https://adb-1234567890123456.7.azuredatabricks.net)' + )); + } + const host = normalizeHost(rawHost); + const meEndpoint = `${host}/api/2.0/preview/scim/v2/Me`; + + const controller = new AbortController(); + const timeout = setTimeout(() => controller.abort(), KEY_VALIDATION_TIMEOUT_MS); + try { + const response = await fetch(meEndpoint, { + method: 'GET', + headers: { + 'Authorization': `Bearer ${apiKey}`, + }, + signal: controller.signal, + }); + + if (response.ok) { + return; + } + + if (response.status === 401 || response.status === 403) { + throw new DatabricksValidationError( + vscode.l10n.t('Invalid Databricks personal access token') + ); + } + + throw new DatabricksValidationError(vscode.l10n.t( + 'Unable to validate Databricks personal access token (HTTP {0})', + String(response.status) + )); + } catch (err) { + if (err instanceof Error && err.name === 'AbortError') { + throw new DatabricksValidationError(vscode.l10n.t( + 'Could not validate Databricks personal access token within {0} seconds', + String(KEY_VALIDATION_TIMEOUT_MS / 1000) + )); + } + if (err instanceof DatabricksValidationError) { + throw err; + } + throw new DatabricksValidationError(vscode.l10n.t( + 'Could not reach the Databricks workspace at {0}. Check the workspace URL and your network connection, then try again.', + host + )); + } finally { + clearTimeout(timeout); + } +} diff --git a/extensions/authentication/src/validation/index.ts b/extensions/authentication/src/validation/index.ts index e35eedd93d41..af189fa41b4f 100644 --- a/extensions/authentication/src/validation/index.ts +++ b/extensions/authentication/src/validation/index.ts @@ -5,6 +5,7 @@ export { validateAnthropicApiKey } from './anthropic'; export { validateCustomProviderApiKey } from './customProvider'; +export { validateDatabricksApiKey } from './databricks'; export { validateDeepSeekApiKey } from './deepseek'; export { normalizeToV1Url, validateFoundryApiKey } from './foundry'; export { validateGeminiApiKey } from './gemini'; diff --git a/product.json b/product.json index ee7102dc1833..8565c03102eb 100644 --- a/product.json +++ b/product.json @@ -747,6 +747,10 @@ "google-cloud": [ "positron.authentication", "posit.assistant" + ], + "databricks": [ + "positron.authentication", + "posit.assistant" ] }, "onboardingKeymaps": [ diff --git a/src/vs/workbench/contrib/positronAssistant/browser/components/languageModelButton.tsx b/src/vs/workbench/contrib/positronAssistant/browser/components/languageModelButton.tsx index d712442803b3..032ca9f06a64 100644 --- a/src/vs/workbench/contrib/positronAssistant/browser/components/languageModelButton.tsx +++ b/src/vs/workbench/contrib/positronAssistant/browser/components/languageModelButton.tsx @@ -9,6 +9,7 @@ import { localize } from '../../../../../nls.js'; import { Button } from '../../../../../base/browser/ui/positronComponents/button/button.js'; import { VerticalStack } from '../../../../browser/positronComponents/positronModalDialog/components/verticalStack.js'; import Claude from '../icons/claude.js'; +import Databricks from '../icons/databricks.js'; import DeepSeek from '../icons/deepseek.js'; import Gemini from '../icons/gemini.js'; import GithubCopilot from '../icons/githubCopilot.js'; @@ -94,6 +95,8 @@ export const LanguageModelIcon = (props: { provider: string; logoUrl?: string }) return ; case 'snowflake-cortex': return ; + case 'databricks': + return ; case 'openai-compatible': return
; case 'error': diff --git a/src/vs/workbench/contrib/positronAssistant/browser/components/languageModelConfigComponent.tsx b/src/vs/workbench/contrib/positronAssistant/browser/components/languageModelConfigComponent.tsx index b5c3fd4fbe9d..2718ff681870 100644 --- a/src/vs/workbench/contrib/positronAssistant/browser/components/languageModelConfigComponent.tsx +++ b/src/vs/workbench/contrib/positronAssistant/browser/components/languageModelConfigComponent.tsx @@ -168,7 +168,7 @@ export const LanguageModelConfigComponent = (props: LanguageModelConfigComponent const hasAutoconfigure = !!source.defaults.autoconfigure && source.defaults.autoconfigure.signedIn && authStatus === AuthStatus.SIGNED_IN; const showApiKeyInput = authMethod === AuthMethod.API_KEY && authStatus !== AuthStatus.SIGNED_IN && !hasAutoconfigure; const showCancelButton = authMethod === AuthMethod.OAUTH && authStatus === AuthStatus.SIGNING_IN && !hasAutoconfigure; - const showBaseUrl = (authMethod === AuthMethod.API_KEY || authMethod === AuthMethod.NONE) && source.supportedOptions?.includes('baseUrl') && !hasAutoconfigure; + const showBaseUrl = (authMethod === AuthMethod.API_KEY || authMethod === AuthMethod.OAUTH || authMethod === AuthMethod.NONE) && source.supportedOptions?.includes('baseUrl') && !hasAutoconfigure; // This currently only updates the API key for the provider, but in the future it may be extended to support // additional configuration options for language models. @@ -176,7 +176,17 @@ export const LanguageModelConfigComponent = (props: LanguageModelConfigComponent props.onChange({ ...props.config, apiKey: newApiKey }); }; + // Under OAuth the base URL (e.g. the Databricks workspace URL) is an input to + // the sign-in flow, and once signed in (either method) it is read-only context + // for the sign-out action — in both cases it renders above the button. Only + // the API-key entry layout (inputs, then button, then base URL) is preserved. + const baseUrlElement = showBaseUrl + ? props.onChange({ ...config, baseUrl: newBaseUrl })} /> + : null; + const showBaseUrlAboveSignIn = authMethod === AuthMethod.OAUTH || authStatus === AuthStatus.SIGNED_IN; + return <> + {showBaseUrlAboveSignIn && baseUrlElement} {!hasAutoconfigure &&
{showApiKeyInput && } props.onSignIn(apiKeyInputRef.current?.value)} /> @@ -187,7 +197,7 @@ export const LanguageModelConfigComponent = (props: LanguageModelConfigComponent }
} {source.provider.id === 'copilot-auth' && authStatus === AuthStatus.SIGNED_IN && } - {showBaseUrl && props.onChange({ ...config, baseUrl: newBaseUrl })} />} + {!showBaseUrlAboveSignIn && baseUrlElement} ; @@ -195,8 +205,28 @@ export const LanguageModelConfigComponent = (props: LanguageModelConfigComponent const DEPLOYMENT_URL_PATTERN = /\/openai\/deployments\//; const SNOWFLAKE_PROVIDER_ID = 'snowflake-cortex'; +const DATABRICKS_PROVIDER_ID = 'databricks'; const BaseUrl = (props: { baseUrl?: string; signedIn?: boolean; onChange: (newBaseUrl: string) => void; provider: IProvider }) => { + // For Databricks, baseUrl holds the workspace URL: relabel as "Workspace URL". + if (props.provider.id === DATABRICKS_PROVIDER_ID) { + const workspaceUrlLabel = localize('positron.languageModelConfig.databricksWorkspaceUrlInputLabel', 'Workspace URL'); + return ( +
+ { + props.signedIn ? +

{localize('positron.languageModelConfig.databricksWorkspaceUrlSignedIn', "Workspace URL: {0}", props.baseUrl)}

+ : + { props.onChange(e.currentTarget.value); }} /> + } +
+ ); + } + // For Snowflake, baseUrl holds the bare account, not a URL: relabel as // "Account Identifier" and pass through. Don't make it a URL input (#13750). if (props.provider.id === SNOWFLAKE_PROVIDER_ID) { diff --git a/src/vs/workbench/contrib/positronAssistant/browser/icons/databricks.svg b/src/vs/workbench/contrib/positronAssistant/browser/icons/databricks.svg new file mode 100644 index 000000000000..2ea276565c87 --- /dev/null +++ b/src/vs/workbench/contrib/positronAssistant/browser/icons/databricks.svg @@ -0,0 +1 @@ +Databricks diff --git a/src/vs/workbench/contrib/positronAssistant/browser/icons/databricks.tsx b/src/vs/workbench/contrib/positronAssistant/browser/icons/databricks.tsx new file mode 100644 index 000000000000..7c09ec55e16c --- /dev/null +++ b/src/vs/workbench/contrib/positronAssistant/browser/icons/databricks.tsx @@ -0,0 +1,26 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (C) 2026 Posit Software, PBC. All rights reserved. + * Licensed under the Elastic License 2.0. See LICENSE.txt for license information. + *--------------------------------------------------------------------------------------------*/ + +import type { SVGProps } from 'react'; + +export const Databricks = (props: SVGProps) => ( + + + +); +export default Databricks; diff --git a/src/vs/workbench/contrib/positronAssistant/browser/languageModelModalDialog.css b/src/vs/workbench/contrib/positronAssistant/browser/languageModelModalDialog.css index aa0b0c633924..312aa98d8390 100644 --- a/src/vs/workbench/contrib/positronAssistant/browser/languageModelModalDialog.css +++ b/src/vs/workbench/contrib/positronAssistant/browser/languageModelModalDialog.css @@ -114,12 +114,37 @@ height: 2rem; border-radius: 5px; border: 1px solid var(--vscode-positronModalDialog-buttonBorder); + color: var(--vscode-positronModalDialog-buttonForeground); background: var(--vscode-positronModalDialog-buttonBackground); width: 80px; min-width: 80px; text-align: center !important; } +.language-model.button.sign-in:hover, +.language-model.button.cancel:hover { + background: var(--vscode-positronModalDialog-buttonHoverBackground); +} + +/* + VS Code 1.109.0 changed secondary buttons to have transparent background and no border in dark + mode. Modal dialog buttons need a visible border, mirroring .action-bar-button:not(.default) + in positronModalDialog.css. +*/ +.vs-dark .language-model.button.sign-in:not(.default), +.vs-dark .language-model.button.cancel { + border: 1px solid var(--vscode-button-border); +} + +.language-model.button.sign-in.default { + color: var(--vscode-positronModalDialog-defaultButtonForeground); + background: var(--vscode-positronModalDialog-defaultButtonBackground); +} + +.language-model.button.sign-in.default:hover { + background: var(--vscode-positronModalDialog-defaultButtonHoverBackground); +} + /** layout the children of this container horizontally */ .language-model-authentication-method-container .radio-group { display: flex; diff --git a/src/vs/workbench/contrib/positronAssistant/browser/languageModelModalDialog.tsx b/src/vs/workbench/contrib/positronAssistant/browser/languageModelModalDialog.tsx index 90ba75752a3b..9681ef389023 100644 --- a/src/vs/workbench/contrib/positronAssistant/browser/languageModelModalDialog.tsx +++ b/src/vs/workbench/contrib/positronAssistant/browser/languageModelModalDialog.tsx @@ -129,6 +129,9 @@ const LanguageModelConfiguration = (props: React.PropsWithChildren(); + // The auth method the user explicitly selected via the radio group, if any. + // Only honored while it is valid for the selected provider. + const [selectedAuthMethod, setSelectedAuthMethod] = useState(); // Ref for preselected provider button to support scrolling into view const preselectedButtonRef = useRef(null); @@ -253,9 +256,13 @@ const LanguageModelConfiguration = (props: React.PropsWithChildren { - // We don't currently support more than one auth method per provider. + // Honor the user's explicit selection when the selected provider supports it. + if (selectedAuthMethod && selectedAuthMethod !== AuthMethod.NONE && + selectedProvider.supportedOptions.includes(selectedAuthMethod)) { + return selectedAuthMethod; + } if (selectedProvider.supportedOptions.includes(AuthMethod.OAUTH)) { return AuthMethod.OAUTH; } else if (selectedProvider.supportedOptions.includes(AuthMethod.API_KEY)) { @@ -269,6 +276,7 @@ const LanguageModelConfiguration = (props: React.PropsWithChildren { - // TODO: it's not currently possible to change the auth method, as each provider only - // supports one auth method at a time. This is a placeholder for future support. + setSelectedAuthMethod(authMethod as AuthMethod); + setErrorMessage(undefined); }} />
diff --git a/src/vs/workbench/contrib/positronAssistant/test/browser/languageModelConfigComponent.vitest.tsx b/src/vs/workbench/contrib/positronAssistant/test/browser/languageModelConfigComponent.vitest.tsx index 16a9e2ed1ef7..28eefb13759a 100644 --- a/src/vs/workbench/contrib/positronAssistant/test/browser/languageModelConfigComponent.vitest.tsx +++ b/src/vs/workbench/contrib/positronAssistant/test/browser/languageModelConfigComponent.vitest.tsx @@ -102,14 +102,17 @@ describe('LanguageModelConfigComponent base-URL input', () => { const ctx = createTestContainer().withReactServices().build(); const rtl = setupRTLRenderer(() => ctx.reactServices); - function renderConfig(supportedOptions: string[]) { + function renderConfig(supportedOptions: string[], options?: { + authMethod?: AuthMethod; + provider?: { id: string; displayName: string }; + }) { rtl.render( { }} config={{ model: '' }} - source={makeSource({ supportedOptions })} + source={makeSource({ supportedOptions, provider: options?.provider })} onCancel={() => { }} onChange={() => { }} onSignIn={() => { }} @@ -128,4 +131,32 @@ describe('LanguageModelConfigComponent base-URL input', () => { expect(screen.queryByLabelText('Base URL')).not.toBeInTheDocument(); }); + + it('shows the base URL input under OAuth when the provider supports baseUrl', () => { + renderConfig(['oauth', 'apiKey', 'baseUrl'], { authMethod: AuthMethod.OAUTH }); + + expect(screen.getByLabelText('Base URL')).toBeInTheDocument(); + }); + + it('hides the base URL input under OAuth when the provider does not support baseUrl', () => { + renderConfig(['oauth'], { authMethod: AuthMethod.OAUTH }); + + expect(screen.queryByLabelText('Base URL')).not.toBeInTheDocument(); + }); + + it('still shows the base URL input under API key auth', () => { + renderConfig(['apiKey', 'baseUrl'], { authMethod: AuthMethod.API_KEY }); + + expect(screen.getByLabelText('Base URL')).toBeInTheDocument(); + }); + + it('labels the Databricks base URL input as Workspace URL', () => { + renderConfig(['oauth', 'apiKey', 'baseUrl'], { + authMethod: AuthMethod.OAUTH, + provider: { id: 'databricks', displayName: 'Databricks' }, + }); + + expect(screen.getByLabelText('Workspace URL')).toBeInTheDocument(); + expect(screen.queryByLabelText('Base URL')).not.toBeInTheDocument(); + }); });