From 96a5a9d5ed18154c1e2bd187aa3b405bb642ac76 Mon Sep 17 00:00:00 2001 From: James Blair Date: Tue, 9 Jun 2026 16:55:19 -0600 Subject: [PATCH 1/5] Add Databricks authentication provider Adds a Databricks auth provider to the authentication extension so Posit Assistant can resolve Databricks credentials via vscode.authentication.getSession('databricks', ...). - OAuth U2M sign-in (authorization code + PKCE) against the built-in databricks-cli public client, with a purpose-built loopback server on the fixed redirect port 8020 and lazy, deduplicated token refresh - Personal access token path via the base AuthProvider machinery, used on remote/web sessions and as an alternative on desktop - New settings: authentication.databricks.credentials (DATABRICKS_HOST), authentication.databricks.customHeaders, and the experimental assistant.provider.databricks.enabled toggle - PAT validation against /api/2.0/preview/scim/v2/Me - Config dialog: implement the auth-method radio selection so providers with multiple supported methods (OAuth + API key) can switch, show the base URL field under OAuth, and label it "Workspace URL" for Databricks; generalize cancelSignIn to the AuthProvider base class - Mocha extension-host tests for the OAuth helpers, loopback server, and provider; vitest coverage for the new base URL rendering rules --- extensions/authentication/package.json | 39 ++ extensions/authentication/package.nls.json | 3 + extensions/authentication/src/authProvider.ts | 7 + extensions/authentication/src/configDialog.ts | 5 +- extensions/authentication/src/constants.ts | 6 + .../src/databricksAuthProvider.ts | 347 ++++++++++++++++++ .../src/databricksAuthServer.ts | 189 ++++++++++ .../authentication/src/databricksOAuth.ts | 161 ++++++++ extensions/authentication/src/extension.ts | 45 +++ .../authentication/src/positOAuthProvider.ts | 2 +- .../authentication/src/providerSources.ts | 28 ++ .../src/test/databricksAuthProvider.test.ts | 212 +++++++++++ .../src/test/databricksAuthServer.test.ts | 139 +++++++ .../src/test/databricksOAuth.test.ts | 246 +++++++++++++ .../src/validation/databricks.ts | 94 +++++ .../authentication/src/validation/index.ts | 1 + .../languageModelConfigComponent.tsx | 22 +- .../browser/languageModelModalDialog.tsx | 16 +- .../languageModelConfigComponent.vitest.tsx | 75 ++++ 19 files changed, 1627 insertions(+), 10 deletions(-) create mode 100644 extensions/authentication/src/databricksAuthProvider.ts create mode 100644 extensions/authentication/src/databricksAuthServer.ts create mode 100644 extensions/authentication/src/databricksOAuth.ts create mode 100644 extensions/authentication/src/test/databricksAuthProvider.test.ts create mode 100644 extensions/authentication/src/test/databricksAuthServer.test.ts create mode 100644 extensions/authentication/src/test/databricksOAuth.test.ts create mode 100644 extensions/authentication/src/validation/databricks.ts diff --git a/extensions/authentication/package.json b/extensions/authentication/package.json index c77be8896863..6382e10d1f60 100644 --- a/extensions/authentication/package.json +++ b/extensions/authentication/package.json @@ -19,6 +19,7 @@ "onAuthenticationRequest:openai-compatible", "onAuthenticationRequest:google", "onAuthenticationRequest:google-cloud", + "onAuthenticationRequest:databricks", "onCommand:authentication.configureProviders", "onCommand:authentication.migrateApiKey" ], @@ -68,6 +69,10 @@ { "id": "google-cloud", "label": "Google Vertex AI" + }, + { + "id": "databricks", + "label": "Databricks" } ], "commands": [ @@ -276,6 +281,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": false, @@ -291,6 +322,14 @@ "tags": [ "experimental" ] + }, + "assistant.provider.databricks.enabled": { + "type": "boolean", + "default": false, + "markdownDescription": "%configuration.provider.databricks.enabled.description%", + "tags": [ + "experimental" + ] } } } diff --git a/extensions/authentication/package.nls.json b/extensions/authentication/package.nls.json index 5ffb33c81dba..1ce73cd12a2a 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 Assistant.\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 Google Vertex AI as a language model provider for Positron Assistant.\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 Assistant.\n\nAuthenticate to the provider in the [language model provider dialog](command:authentication.configureProviders?%5B%7B%22preselectedProviderId%22%3A%22deepseek-api%22%7D%5D).", "commands.configureProviders.title": "Configure Language Model Providers", diff --git a/extensions/authentication/src/authProvider.ts b/extensions/authentication/src/authProvider.ts index 71868f8b90a7..dfeb15e8be00 100644 --- a/extensions/authentication/src/authProvider.ts +++ b/extensions/authentication/src/authProvider.ts @@ -75,6 +75,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 baad4b042538..8401ebef836e 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'; @@ -206,9 +205,7 @@ export async function showConfigurationDialog( } case 'cancel': { const provider = authProviders.get(config.provider); - 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 cb2647072291..39c161c3ffe8 100644 --- a/extensions/authentication/src/constants.ts +++ b/extensions/authentication/src/constants.ts @@ -21,3 +21,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 e4fcebc5160b..4adff2a32fcc 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, @@ -34,6 +36,8 @@ import { import { FOUNDRY_MANAGED_CREDENTIALS, hasManagedCredentials } from './managedCredentials'; import { detectSnowflakeCredentials, getSnowflakeConnectionsTomlPath } from './snowflakeCredentials'; 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'; @@ -62,6 +66,7 @@ export async function activate(context: vscode.ExtensionContext) { await registerGeminiProvider(context); await registerGoogleVertexProvider(context); await registerDeepSeekProvider(context); + registerDatabricksProvider(context); registerCustomProvider(context); // Register provider metadata so the Settings UI shows per-provider @@ -648,6 +653,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 76df5f3e934a..2e0a8eda3703 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 64482e202ffa..835bd92f63df 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, @@ -86,6 +87,11 @@ export const PROVIDER_METADATA: Record = { displayName: 'DeepSeek', settingName: 'deepseek', }, + databricks: { + id: DATABRICKS_AUTH_PROVIDER_ID, + displayName: 'Databricks', + settingName: 'databricks', + }, }; export function getProviderSources(): positron.ai.LanguageModelSource[] { @@ -95,6 +101,14 @@ export function getProviderSources(): positron.ai.LanguageModelSource[] { const vertexFromEnv = !!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, @@ -250,5 +264,19 @@ 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: { + name: 'Databricks', + 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 f3f3f66fa3ec..d24cbba530c3 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/src/vs/workbench/contrib/positronAssistant/browser/components/languageModelConfigComponent.tsx b/src/vs/workbench/contrib/positronAssistant/browser/components/languageModelConfigComponent.tsx index cad002844131..7723b2ede1cd 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 && source.supportedOptions?.includes('baseUrl') && !hasAutoconfigure; + const showBaseUrl = (authMethod === AuthMethod.API_KEY || authMethod === AuthMethod.OAUTH) && 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. @@ -195,8 +195,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/languageModelModalDialog.tsx b/src/vs/workbench/contrib/positronAssistant/browser/languageModelModalDialog.tsx index b3d329beaddd..1425e7e77cb7 100644 --- a/src/vs/workbench/contrib/positronAssistant/browser/languageModelModalDialog.tsx +++ b/src/vs/workbench/contrib/positronAssistant/browser/languageModelModalDialog.tsx @@ -114,6 +114,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); @@ -243,9 +246,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)) { @@ -258,6 +265,7 @@ const LanguageModelConfiguration = (props: React.PropsWithChildren { setSelectedProvider(provider); setProviderConfig(providerSourceToConfig(provider)); + setSelectedAuthMethod(undefined); setShowProgress(false); setErrorMessage(undefined); }; @@ -420,8 +428,8 @@ 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 a1af3f848408..341b19b74d22 100644 --- a/src/vs/workbench/contrib/positronAssistant/test/browser/languageModelConfigComponent.vitest.tsx +++ b/src/vs/workbench/contrib/positronAssistant/test/browser/languageModelConfigComponent.vitest.tsx @@ -93,3 +93,78 @@ describe('LanguageModelConfigComponent ProviderNotice', () => { }); }); }); + +describe('LanguageModelConfigComponent BaseUrl', () => { + const ctx = createTestContainer().withReactServices().build(); + const rtl = setupRTLRenderer(() => ctx.reactServices); + + function renderComponent(options: { + providerId: string; + authMethod: AuthMethod; + supportedOptions: IPositronLanguageModelSource['supportedOptions']; + }): void { + const source: IPositronLanguageModelSource = { + type: PositronLanguageModelType.Chat, + provider: { + id: options.providerId, + displayName: options.providerId, + settingName: options.providerId, + }, + supportedOptions: options.supportedOptions, + defaults: { name: '', model: '' }, + }; + rtl.render( + { }} + config={{ type: PositronLanguageModelType.Chat, provider: options.providerId, name: '', model: '' }} + source={source} + onCancel={() => { }} + onChange={() => { }} + onSignIn={() => { }} + /> + ); + } + + it('shows the base URL input under OAuth when the provider supports baseUrl', () => { + renderComponent({ + providerId: 'some-provider', + authMethod: AuthMethod.OAUTH, + supportedOptions: ['oauth', 'apiKey', 'baseUrl'], + }); + + expect(screen.getByLabelText('Base URL')).toBeInTheDocument(); + }); + + it('hides the base URL input under OAuth when the provider does not support baseUrl', () => { + renderComponent({ + providerId: 'some-provider', + authMethod: AuthMethod.OAUTH, + supportedOptions: ['oauth'], + }); + + expect(screen.queryByLabelText('Base URL')).not.toBeInTheDocument(); + }); + + it('still shows the base URL input under API key auth', () => { + renderComponent({ + providerId: 'some-provider', + authMethod: AuthMethod.API_KEY, + supportedOptions: ['apiKey', 'baseUrl'], + }); + + expect(screen.getByLabelText('Base URL')).toBeInTheDocument(); + }); + + it('labels the Databricks base URL input as Workspace URL', () => { + renderComponent({ + providerId: 'databricks', + authMethod: AuthMethod.OAUTH, + supportedOptions: ['oauth', 'apiKey', 'baseUrl'], + }); + + expect(screen.getByLabelText('Workspace URL')).toBeInTheDocument(); + expect(screen.queryByLabelText('Base URL')).not.toBeInTheDocument(); + }); +}); From 070165aa19f04017e289a0fa220b24d54189be7d Mon Sep 17 00:00:00 2001 From: James Blair Date: Tue, 9 Jun 2026 19:39:40 -0600 Subject: [PATCH 2/5] Add Databricks logo to the language model provider button Uses the official Databricks mark (path data from Simple Icons) in brand red, rendered for the 'databricks' auth provider id in the provider config dialog, matching the other branded providers. --- .../components/languageModelButton.tsx | 3 +++ .../browser/icons/databricks.svg | 1 + .../browser/icons/databricks.tsx | 26 +++++++++++++++++++ 3 files changed, 30 insertions(+) create mode 100644 src/vs/workbench/contrib/positronAssistant/browser/icons/databricks.svg create mode 100644 src/vs/workbench/contrib/positronAssistant/browser/icons/databricks.tsx diff --git a/src/vs/workbench/contrib/positronAssistant/browser/components/languageModelButton.tsx b/src/vs/workbench/contrib/positronAssistant/browser/components/languageModelButton.tsx index 0ebd9b06d197..981aca7c1157 100644 --- a/src/vs/workbench/contrib/positronAssistant/browser/components/languageModelButton.tsx +++ b/src/vs/workbench/contrib/positronAssistant/browser/components/languageModelButton.tsx @@ -8,6 +8,7 @@ import * as React from 'react'; 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'; @@ -74,6 +75,8 @@ export const LanguageModelIcon = (props: { provider: string }) => { return ; case 'snowflake-cortex': return ; + case 'databricks': + return ; case 'openai-compatible': return
; case 'error': 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; From 9b6bd371b964b7e678e6996cb9e28612ab3cce5b Mon Sep 17 00:00:00 2001 From: James Blair Date: Tue, 9 Jun 2026 21:57:26 -0600 Subject: [PATCH 3/5] Grant Posit Assistant trusted access to Databricks auth sessions Without a trustedExtensionAuthAccess entry, posit.assistant's silent getSession('databricks') returns no session even when signed in, so Databricks models never appear. --- product.json | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/product.json b/product.json index 579183e7ba03..10a16b120c14 100644 --- a/product.json +++ b/product.json @@ -750,6 +750,10 @@ "google-cloud": [ "positron.authentication", "posit.assistant" + ], + "databricks": [ + "positron.authentication", + "posit.assistant" ] }, "onboardingKeymaps": [ From 5f7555ebc582e31364eeb8294a8537ace4d44fad Mon Sep 17 00:00:00 2001 From: James Blair Date: Wed, 10 Jun 2026 16:01:36 -0600 Subject: [PATCH 4/5] Fix sign-in button theming and OAuth field order in provider config dialog The bespoke sign-in/cancel button styles predate VS Code 1.109's change to transparent borderless secondary buttons in dark mode, leaving the button nearly invisible there; mirror the canonical .action-bar-button rules (foreground, hover, dark-mode border, default variant). Under OAuth, the base URL (e.g. the Databricks workspace URL) is an input to the sign-in flow, so render it above the sign-in button. --- .../languageModelConfigComponent.tsx | 11 +++++++- .../browser/languageModelModalDialog.css | 25 +++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/src/vs/workbench/contrib/positronAssistant/browser/components/languageModelConfigComponent.tsx b/src/vs/workbench/contrib/positronAssistant/browser/components/languageModelConfigComponent.tsx index 7723b2ede1cd..5844be941b5f 100644 --- a/src/vs/workbench/contrib/positronAssistant/browser/components/languageModelConfigComponent.tsx +++ b/src/vs/workbench/contrib/positronAssistant/browser/components/languageModelConfigComponent.tsx @@ -176,7 +176,16 @@ 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, so it renders above the sign-in button; under API key the + // existing layout (inputs, then button, then base URL) is preserved. + const baseUrlElement = showBaseUrl + ? props.onChange({ ...config, baseUrl: newBaseUrl })} /> + : null; + const showBaseUrlAboveSignIn = authMethod === AuthMethod.OAUTH; + return <> + {showBaseUrlAboveSignIn && baseUrlElement} {!hasAutoconfigure &&
{showApiKeyInput && } props.onSignIn(apiKeyInputRef.current?.value)} /> @@ -187,7 +196,7 @@ export const LanguageModelConfigComponent = (props: LanguageModelConfigComponent }
} {source.provider.id === 'copilot-auth' && authStatus === AuthStatus.SIGNED_IN && } - {showBaseUrl && props.onChange({ ...config, baseUrl: newBaseUrl })} />} + {!showBaseUrlAboveSignIn && baseUrlElement} ; diff --git a/src/vs/workbench/contrib/positronAssistant/browser/languageModelModalDialog.css b/src/vs/workbench/contrib/positronAssistant/browser/languageModelModalDialog.css index b6ece73c33bd..526df7eded14 100644 --- a/src/vs/workbench/contrib/positronAssistant/browser/languageModelModalDialog.css +++ b/src/vs/workbench/contrib/positronAssistant/browser/languageModelModalDialog.css @@ -103,12 +103,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; From 8c251eb7318ed583e5ce53035c27efb833550b9e Mon Sep 17 00:00:00 2001 From: James Blair Date: Wed, 10 Jun 2026 17:09:25 -0600 Subject: [PATCH 5/5] Render base URL above the sign-out button when signed in Extends the OAuth field-order fix to the signed-in state of both auth methods: the read-only base URL (e.g. Databricks workspace URL, Snowflake account) is context for the sign-out action, so it renders above the button. API-key entry layout is unchanged. --- .../browser/components/languageModelConfigComponent.tsx | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/vs/workbench/contrib/positronAssistant/browser/components/languageModelConfigComponent.tsx b/src/vs/workbench/contrib/positronAssistant/browser/components/languageModelConfigComponent.tsx index 5844be941b5f..d56636e96ea6 100644 --- a/src/vs/workbench/contrib/positronAssistant/browser/components/languageModelConfigComponent.tsx +++ b/src/vs/workbench/contrib/positronAssistant/browser/components/languageModelConfigComponent.tsx @@ -177,12 +177,13 @@ export const LanguageModelConfigComponent = (props: LanguageModelConfigComponent }; // Under OAuth the base URL (e.g. the Databricks workspace URL) is an input to - // the sign-in flow, so it renders above the sign-in button; under API key the - // existing layout (inputs, then button, then base URL) is preserved. + // 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; + const showBaseUrlAboveSignIn = authMethod === AuthMethod.OAUTH || authStatus === AuthStatus.SIGNED_IN; return <> {showBaseUrlAboveSignIn && baseUrlElement}