diff --git a/.gitignore b/.gitignore index f4e2c6d..a79a473 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ node_modules/ dist/ *.tsbuildinfo +*.tgz diff --git a/memory-bank/aiCredentialStore.md b/memory-bank/aiCredentialStore.md index 34af9b8..10e54ab 100644 --- a/memory-bank/aiCredentialStore.md +++ b/memory-bank/aiCredentialStore.md @@ -8,27 +8,26 @@ package: ai-credentials ## Overview -`ai-credentials` provides browser-safe credential types and shaping logic, plus -a generic typed key-value secret store for disk persistence. It serves as the -single source of truth for credential interfaces (`ProviderCredentials`, -`ApiKeyCredentials`, etc.) and the `shapeCredentials()` function that transforms -raw auth tokens into provider-ready credential objects. - -It is a **leaf** package: it imports nothing from `ai-config`, -`ai-provider-bridge`, or any host application, and the two sibling packages do -not import it (they re-export from it). Its only runtime dependencies are -`chokidar` (watching) and `proper-lockfile` (cross-process locking), used by the -`/store` entrypoint. The `/types` entrypoint is fully browser-safe with zero -platform dependencies. +`ai-credentials` provides browser-safe credential types and shaping logic, a +provider-neutral OAuth acquisition controller, a generic typed key-value secret +store, and a credential-aware store backend. It is the single source of truth +for runtime credential interfaces and for the generation protocol that prevents +stale OAuth completions from resurrecting replaced credentials. + +It imports nothing from `ai-config`, `ai-provider-bridge`, or any host +application. The pure root and `/types` entrypoints remain browser-safe; +platform dependencies are isolated to `/store`, `/store-backend`, and +`/positron`. ## Entrypoints -| Entrypoint | Purpose | Browser-safe? | -| ------------------------- | --------------------------------------------------------------------------------- | ------------- | -| `ai-credentials` | Stub root entry (Phase 4: CredentialProvider interface + factory) | Yes | -| `ai-credentials/types` | Credential interfaces, `shapeCredentials()`, `AuthProviderMapping`, `Logger` | **Yes** | -| `ai-credentials/store` | `SingleFileStore` class, `createDefaultStore`, `getDefaultStorePath`, store types | No (Node FS) | -| `ai-credentials/positron` | Stub entry (Phase 4: vscode.authentication backend) | N/A | +| Entrypoint | Purpose | Browser-safe? | +| ------------------------------ | --------------------------------------------------------------------------------- | ------------- | +| `ai-credentials` | CredentialProvider factory, unified acquisition controller, OAuth helpers | Yes | +| `ai-credentials/types` | Credential interfaces, `shapeCredentials()`, `AuthProviderMapping`, `Logger` | **Yes** | +| `ai-credentials/store` | `SingleFileStore` class, `createDefaultStore`, `getDefaultStorePath`, store types | No (Node FS) | +| `ai-credentials/store-backend` | Credential-aware store/env resolver, disk schema, transactions | No (Node FS) | +| `ai-credentials/positron` | `vscode.authentication` backend | No (VS Code) | ### `/types` — Browser-safe credential types and shaping @@ -53,10 +52,10 @@ The store entrypoint provides a typed key-value store backed by a single JSON file on disk. It is built for small amounts of sensitive data (credentials, OAuth tokens, API keys). -**The store owns where and how bytes hit disk; it does not own credential -meaning.** OAuth semantics, provider grouping, and auth status stay with the -caller. Values are fully generic — callers supply their own types via method -type parameters. +**The generic `/store` entrypoint owns where and how bytes hit disk; it does not +own credential meaning.** The separate `/store-backend` entrypoint owns tolerant +credential parsing, source normalization, environment precedence, status, and +compare-and-commit OAuth transactions. Values in `/store` remain fully generic. ```ts import { createDefaultStore, getDefaultStorePath } from "ai-credentials/store"; @@ -178,8 +177,10 @@ directories (`0o700`) and an empty `{}` file (`0o600`). | `src/store/SingleFileStore.ts` | The store: read/write/lock/watch + private helpers (`writeStore`, `readStore`, `withWriteLock`, `ensureFileExists`, `ensurePermissions`) | | `src/store/defaults.ts` | `getDefaultStorePath()` and `createDefaultStore()` — canonical default path convention | | `src/store/index.ts` | `/store` entrypoint exports | -| `src/index.ts` | Root entrypoint (stub for Phase 4) | -| `src/positron/index.ts` | `/positron` entrypoint (stub for Phase 4) | +| `src/index.ts` | Root CredentialProvider/acquisition entrypoint | +| `src/acquisition.ts` | Unified device-code, PKCE, M2M, refresh, cancellation, and awaited disposal controller | +| `src/store-backend/` | Credential disk schema, store/env resolution, normalization, generation transactions, and status | +| `src/positron/index.ts` | `/positron` vscode.authentication backend | ## Invariants & Design Decisions @@ -200,3 +201,6 @@ directories (`0o700`) and an empty `{}` file (`0o600`). - **Graceful degradation** — invalid JSON and watcher-init failures log and recover rather than throw. - **Disposable watch lifecycle** — mirrors VS Code's resource pattern. +- **One acquisition controller per provider handle** — generalized store-backed + backends instantiate only `AcquisitionEngine`; `OAuthEngine` exists solely as + a fallback for older backends without generalized acquisition hooks. diff --git a/packages/ai-config/providers.schema.json b/packages/ai-config/providers.schema.json index ddd053d..f62ebe7 100644 --- a/packages/ai-config/providers.schema.json +++ b/packages/ai-config/providers.schema.json @@ -2909,6 +2909,241 @@ }, "type": "object" }, + "databricks": { + "additionalProperties": false, + "properties": { + "baseUrl": { + "type": "string" + }, + "customHeaders": { + "additionalProperties": { + "type": "string" + }, + "propertyNames": { + "type": "string" + }, + "type": "object" + }, + "databricks": { + "additionalProperties": false, + "properties": { + "host": { + "type": "string" + } + }, + "type": "object" + }, + "enabled": { + "type": "boolean" + }, + "endpoint": { + "type": "string" + }, + "endpoints": { + "additionalProperties": { + "type": "string" + }, + "propertyNames": { + "enum": [ + "anthropic-messages", + "openai-chat", + "openai-responses", + "bedrock-converse", + "google-generative" + ], + "type": "string" + }, + "type": "object" + }, + "models": { + "additionalProperties": false, + "properties": { + "allow": { + "items": { + "type": "string" + }, + "type": "array" + }, + "custom": { + "items": { + "additionalProperties": false, + "properties": { + "baseUrl": { + "type": "string" + }, + "family": { + "type": "string" + }, + "id": { + "minLength": 1, + "type": "string" + }, + "maxContextLength": { + "exclusiveMinimum": 0, + "maximum": 9007199254740991, + "type": "integer" + }, + "maxInputTokens": { + "exclusiveMinimum": 0, + "maximum": 9007199254740991, + "type": "integer" + }, + "maxOutputTokens": { + "exclusiveMinimum": 0, + "maximum": 9007199254740991, + "type": "integer" + }, + "name": { + "minLength": 1, + "type": "string" + }, + "protocol": { + "enum": [ + "anthropic-messages", + "openai-chat", + "openai-responses", + "bedrock-converse", + "google-generative" + ], + "type": "string" + }, + "supportedInputMediaTypes": { + "items": { + "type": "string" + }, + "type": "array" + }, + "supportsImages": { + "type": "boolean" + }, + "supportsToolResultImages": { + "type": "boolean" + }, + "supportsTools": { + "type": "boolean" + }, + "supportsWebSearch": { + "type": "boolean" + }, + "thinkingEffortLevels": { + "items": { + "type": "string" + }, + "type": "array" + } + }, + "required": [ + "id", + "name", + "maxContextLength", + "supportsTools", + "supportsImages", + "supportsToolResultImages", + "supportsWebSearch" + ], + "type": "object" + }, + "type": "array" + }, + "deny": { + "items": { + "type": "string" + }, + "type": "array" + }, + "discovery": { + "enum": [ + "auto", + "off" + ], + "type": "string" + }, + "overrides": { + "additionalProperties": { + "additionalProperties": false, + "properties": { + "baseUrl": { + "type": "string" + }, + "family": { + "type": "string" + }, + "maxContextLength": { + "exclusiveMinimum": 0, + "maximum": 9007199254740991, + "type": "integer" + }, + "maxInputTokens": { + "exclusiveMinimum": 0, + "maximum": 9007199254740991, + "type": "integer" + }, + "maxOutputTokens": { + "exclusiveMinimum": 0, + "maximum": 9007199254740991, + "type": "integer" + }, + "name": { + "type": "string" + }, + "protocol": { + "enum": [ + "anthropic-messages", + "openai-chat", + "openai-responses", + "bedrock-converse", + "google-generative" + ], + "type": "string" + }, + "supportedInputMediaTypes": { + "items": { + "type": "string" + }, + "type": "array" + }, + "supportsImages": { + "type": "boolean" + }, + "supportsToolResultImages": { + "type": "boolean" + }, + "supportsTools": { + "type": "boolean" + }, + "supportsWebSearch": { + "type": "boolean" + }, + "thinkingEffortLevels": { + "items": { + "type": "string" + }, + "type": "array" + } + }, + "type": "object" + }, + "propertyNames": { + "type": "string" + }, + "type": "object" + } + }, + "type": "object" + }, + "protocol": { + "enum": [ + "anthropic-messages", + "openai-chat", + "openai-responses", + "bedrock-converse", + "google-generative" + ], + "type": "string" + } + }, + "type": "object" + }, "deepseek": { "additionalProperties": false, "properties": { diff --git a/packages/ai-config/src/__tests__/load-config.test.ts b/packages/ai-config/src/__tests__/load-config.test.ts index 041dc33..896c417 100644 --- a/packages/ai-config/src/__tests__/load-config.test.ts +++ b/packages/ai-config/src/__tests__/load-config.test.ts @@ -12,6 +12,7 @@ import { loadResolvedProviderCatalog } from "../node/load-catalog.js"; import { mutateProvidersConfig } from "../node/mutate-config.js"; import type { ProvidersConfig, ResolvedProvider } from "../types.js"; import type { PlatformBaseline } from "../types.js"; +import { BUILTIN_PROVIDER_IDS } from "../vocabulary.js"; // --------------------------------------------------------------------------- // Helpers @@ -72,8 +73,7 @@ describe("loadResolvedProviderCatalog", () => { logger: mockLogger, }); - // 14 built-in providers - expect(catalog.length).toBe(14); + expect(catalog.length).toBe(BUILTIN_PROVIDER_IDS.length); expect(findProvider(catalog, "positai")?.enabled).toBe(true); expect(findProvider(catalog, "anthropic")?.enabled).toBe(true); }); @@ -87,7 +87,7 @@ describe("loadResolvedProviderCatalog", () => { logger: mockLogger, }); - expect(catalog.length).toBe(14); + expect(catalog.length).toBe(BUILTIN_PROVIDER_IDS.length); }); it("should apply RStudio baseline (positai only)", async () => { @@ -216,7 +216,7 @@ describe("loadResolvedProviderCatalog", () => { logger: mockLogger, }); - expect(catalog.length).toBe(14); + expect(catalog.length).toBe(BUILTIN_PROVIDER_IDS.length); expect(mockLogger.warn).toHaveBeenCalledWith( expect.stringContaining("Failed to parse TEST_ENFORCED"), ); @@ -291,7 +291,7 @@ describe("loadResolvedProviderCatalog", () => { }); // Should fall back to user config (no custom provider) - expect(catalog.length).toBe(14); + expect(catalog.length).toBe(BUILTIN_PROVIDER_IDS.length); expect(findProvider(catalog, "env-only-gateway")).toBeUndefined(); expect(mockLogger.warn).toHaveBeenCalledWith( expect.stringContaining("invalid merged result"), @@ -474,8 +474,8 @@ describe("loadResolvedProviderCatalog", () => { logger: mockLogger, }); - // 14 built-ins + 1 custom - expect(catalog.length).toBe(15); + // built-ins + 1 custom + expect(catalog.length).toBe(BUILTIN_PROVIDER_IDS.length + 1); const custom = findProvider(catalog, "my-gateway"); expect(custom).toBeDefined(); @@ -558,7 +558,7 @@ describe("loadResolvedProviderCatalog", () => { logger: mockLogger, }); - expect(catalog.length).toBe(14); // defaults + expect(catalog.length).toBe(BUILTIN_PROVIDER_IDS.length); // defaults expect(mockLogger.warn).toHaveBeenCalledWith(expect.stringContaining("Failed to parse")); }); @@ -572,7 +572,7 @@ describe("loadResolvedProviderCatalog", () => { logger: mockLogger, }); - expect(catalog.length).toBe(14); // defaults + expect(catalog.length).toBe(BUILTIN_PROVIDER_IDS.length); // defaults expect(mockLogger.warn).toHaveBeenCalledWith(expect.stringContaining("Validation errors")); }); @@ -591,7 +591,7 @@ describe("loadResolvedProviderCatalog", () => { logger: mockLogger, }); - expect(catalog.length).toBe(14); // defaults + expect(catalog.length).toBe(BUILTIN_PROVIDER_IDS.length); // defaults expect(findProvider(catalog, "anthropic")?.connection.aws).toBeUndefined(); expect(mockLogger.warn).toHaveBeenCalledWith(expect.stringContaining("Validation errors")); }); @@ -652,7 +652,7 @@ describe("loadResolvedProviderCatalog", () => { additionalSources: [emptyHostSource], }); - expect(catalog.length).toBe(14); + expect(catalog.length).toBe(BUILTIN_PROVIDER_IDS.length); expect(findProvider(catalog, "anthropic")?.connection.baseUrl).toBeUndefined(); }); diff --git a/packages/ai-config/src/__tests__/positron-source.test.ts b/packages/ai-config/src/__tests__/positron-source.test.ts index 35459f5..4dcdbeb 100644 --- a/packages/ai-config/src/__tests__/positron-source.test.ts +++ b/packages/ai-config/src/__tests__/positron-source.test.ts @@ -22,6 +22,7 @@ function fakeReader( customHeaders?: Record>; awsRegion?: string; snowflake?: { host?: string; account?: string; home?: string }; + databricks?: { host?: string }; } = {}, ): PositronAuthSettingReader { return { @@ -29,6 +30,7 @@ function fakeReader( getCustomHeaders: (configKey) => overrides.customHeaders?.[configKey], getAwsRegion: () => overrides.awsRegion, getSnowflake: () => overrides.snowflake, + getDatabricks: () => overrides.databricks, }; } @@ -53,6 +55,11 @@ const SNOWFLAKE: PositronAuthSettingDescriptor = { configKey: "snowflake", read: "snowflake", }; +const DATABRICKS: PositronAuthSettingDescriptor = { + providerId: "databricks", + configKey: "databricks", + read: "databricks", +}; describe("buildAuthenticationFragment", () => { it("returns an empty fragment when nothing is set", () => { @@ -156,6 +163,30 @@ describe("buildAuthenticationFragment", () => { }); }); + it("emits the databricks host under the databricks section (+ customHeaders), not baseUrl", () => { + const reader = fakeReader({ + databricks: { host: "https://adb-123.4.azuredatabricks.net" }, + customHeaders: { databricks: { "x-databricks-use-coding-agent-mode": "true" } }, + }); + const fragment = buildAuthenticationFragment(reader, [DATABRICKS]); + // The host must NOT be emitted as baseUrl: per-model endpoint resolution + // falls back to the provider baseUrl, which would route chat requests to + // the bare workspace host instead of the serving-endpoints path. + expect(fragment).toEqual({ + providers: { + databricks: { + databricks: { host: "https://adb-123.4.azuredatabricks.net" }, + customHeaders: { "x-databricks-use-coding-agent-mode": "true" }, + }, + }, + }); + }); + + it("omits databricks when no host is set", () => { + const fragment = buildAuthenticationFragment(fakeReader(), [DATABRICKS]); + expect(fragment).toEqual({}); + }); + it("applies the descriptor's normalizeBaseUrl to a set base URL", () => { const normalizeBaseUrl = (url: string) => url === "https://bare.example" ? "https://bare.example/v1" : url; diff --git a/packages/ai-config/src/__tests__/resolve-catalog.test.ts b/packages/ai-config/src/__tests__/resolve-catalog.test.ts index 1425f05..86999cb 100644 --- a/packages/ai-config/src/__tests__/resolve-catalog.test.ts +++ b/packages/ai-config/src/__tests__/resolve-catalog.test.ts @@ -138,8 +138,8 @@ describe("resolveProviderCatalog — invalid merge tolerance", () => { logger, }); - // 14 built-ins, no custom provider leaked in. - expect(catalog.length).toBe(14); + // 15 built-ins, no custom provider leaked in. + expect(catalog.length).toBe(15); expect(find(catalog, "ghost-gw")).toBeUndefined(); expect(find(catalog, "anthropic")?.enabled).toBe(true); expect(logger.warn).toHaveBeenCalledWith(expect.stringContaining("invalid merged result")); @@ -186,7 +186,7 @@ describe("resolveProviderCatalog — tightened-schema recovery", () => { }); // The valid user source still resolves; the forbidden overlay is gone. - expect(catalog.length).toBe(14); + expect(catalog.length).toBe(15); expect(find(catalog, "openai")?.enabled).toBe(true); expect(find(catalog, "anthropic")?.connection.aws).toBeUndefined(); expect(logger.warn).toHaveBeenCalledWith(expect.stringContaining("invalid merged result")); diff --git a/packages/ai-config/src/__tests__/watch-catalog.test.ts b/packages/ai-config/src/__tests__/watch-catalog.test.ts index 3de594a..6403ce4 100644 --- a/packages/ai-config/src/__tests__/watch-catalog.test.ts +++ b/packages/ai-config/src/__tests__/watch-catalog.test.ts @@ -241,7 +241,7 @@ describe("watchResolvedProviderCatalog", () => { if (changes.length > 0) { const lastChange = changes[changes.length - 1]; - expect(lastChange.catalog.length).toBe(14); // all built-ins + expect(lastChange.catalog.length).toBe(15); // all built-ins } }); }); diff --git a/packages/ai-config/src/build-catalog.ts b/packages/ai-config/src/build-catalog.ts index 3365769..1cce43d 100644 --- a/packages/ai-config/src/build-catalog.ts +++ b/packages/ai-config/src/build-catalog.ts @@ -46,6 +46,7 @@ interface ConnectionEnvMapping { positaiLogin?: { host?: string; clientId?: string; scope?: string }; aws?: { region?: string; profile?: string }; googleCloud?: { project?: string; location?: string }; + databricks?: { host?: string }; } const CONNECTION_ENV_MAPPINGS: Readonly> = { @@ -74,6 +75,10 @@ const CONNECTION_ENV_MAPPINGS: Readonly> = "ms-foundry": { baseUrl: "MS_FOUNDRY_BASE_URL" }, "snowflake-cortex": { baseUrl: "SNOWFLAKE_BASE_URL" }, deepseek: { baseUrl: "DEEPSEEK_BASE_URL" }, + // The standard Databricks CLI/SDK variable. Maps into the `databricks` + // section (NOT baseUrl): the workspace host is not a chat base URL — the + // bridge derives the serving-endpoints / AI Gateway URL from it. + databricks: { databricks: { host: "DATABRICKS_HOST" } }, }; // --------------------------------------------------------------------------- @@ -104,6 +109,7 @@ const BUILTIN_CLIENT_KIND = { "snowflake-cortex": "snowflake", "ms-foundry": "ms-foundry", deepseek: "deepseek", + databricks: "databricks", } as const satisfies Record; /** @@ -222,6 +228,7 @@ function resolveConnection( aws: mergeOptionalSection(defaults.aws, fromBlock.aws), googleCloud: mergeOptionalSection(defaults.googleCloud, fromBlock.googleCloud), snowflake: mergeOptionalSection(defaults.snowflake, fromBlock.snowflake), + databricks: mergeOptionalSection(defaults.databricks, fromBlock.databricks), }; } @@ -253,6 +260,7 @@ function resolveConnectionFromBlock( aws: superset.aws, googleCloud: superset.googleCloud, snowflake: superset.snowflake, + databricks: superset.databricks, }; } @@ -340,6 +348,14 @@ function applyEnvOverlay( changed = true; } } + if (mapping.databricks) { + const overlay = readEnvSection(mapping.databricks, envVars); + if (overlay) { + result = changed ? result : { ...result }; + result.databricks = result.databricks ? { ...result.databricks, ...overlay } : overlay; + changed = true; + } + } return result; } diff --git a/packages/ai-config/src/positron/authentication-fragment.ts b/packages/ai-config/src/positron/authentication-fragment.ts index c9e280d..4cc3297 100644 --- a/packages/ai-config/src/positron/authentication-fragment.ts +++ b/packages/ai-config/src/positron/authentication-fragment.ts @@ -42,6 +42,11 @@ export interface PositronAuthSettingReader { * SNOWFLAKE_ACCOUNT,SNOWFLAKE_HOME}`, `process.env` fallback). */ getSnowflake(): { host?: string; account?: string; home?: string } | undefined; + /** + * Databricks workspace host (`authentication.databricks.credentials.DATABRICKS_HOST`, + * `process.env` fallback). + */ + getDatabricks(): { host?: string } | undefined; } /** @@ -61,8 +66,13 @@ export interface PositronAuthSettingDescriptor { * - `"aws-region"`: reads `authentication.aws.credentials.AWS_REGION`. * - `"snowflake"`: reads `snowflake.credentials.{SNOWFLAKE_HOST,SNOWFLAKE_ACCOUNT, * SNOWFLAKE_HOME}` (+ `snowflake.customHeaders`), with `process.env` fallback. + * - `"databricks"`: reads `databricks.credentials.DATABRICKS_HOST` + * (+ `databricks.customHeaders`), with `process.env` fallback for the host. + * The host is emitted as the `databricks` connection section — NOT as + * `baseUrl`, which would be picked up by per-model endpoint resolution and + * route chat to the bare workspace host. */ - readonly read: "api-key-connection" | "aws-region" | "snowflake"; + readonly read: "api-key-connection" | "aws-region" | "snowflake" | "databricks"; /** * Optional correction applied to the raw `baseUrl` setting value before it * enters the fragment (only meaningful for `"api-key-connection"` reads). @@ -125,6 +135,8 @@ function buildBlock( return buildAwsRegionBlock(reader); case "snowflake": return buildSnowflakeBlock(reader, descriptor.configKey); + case "databricks": + return buildDatabricksBlock(reader, descriptor.configKey); } } @@ -182,6 +194,27 @@ function buildSnowflakeBlock( return hasKeys(block) ? block : undefined; } +function buildDatabricksBlock( + reader: PositronAuthSettingReader, + configKey: string, +): BuiltinProviderBlock | undefined { + const block: BuiltinProviderBlock = {}; + + // The host is stored raw (like snowflake's) — the credential shaper and the + // bridge's Databricks provider normalize scheme/trailing-slash at use. + const host = reader.getDatabricks()?.host || undefined; + if (host) { + block.databricks = { host }; + } + + const customHeaders = normalizeHeaders(reader.getCustomHeaders(configKey)); + if (customHeaders) { + block.customHeaders = customHeaders; + } + + return hasKeys(block) ? block : undefined; +} + /** Empty header maps normalize to `undefined` (match the shaper's pipeline). */ function normalizeHeaders( headers: Record | undefined, diff --git a/packages/ai-config/src/positron/index.ts b/packages/ai-config/src/positron/index.ts index 5189a5b..e71ba03 100644 --- a/packages/ai-config/src/positron/index.ts +++ b/packages/ai-config/src/positron/index.ts @@ -62,6 +62,14 @@ function createVscodeAuthReader(): PositronAuthSettingReader { home: snowflakeConfig?.SNOWFLAKE_HOME || process.env.SNOWFLAKE_HOME, }; }, + getDatabricks: () => { + const databricksConfig = vscode.workspace + .getConfiguration("authentication.databricks") + .get>("credentials"); + return { + host: databricksConfig?.DATABRICKS_HOST || process.env.DATABRICKS_HOST, + }; + }, }; } diff --git a/packages/ai-config/src/schema.ts b/packages/ai-config/src/schema.ts index 12ae1ae..081988a 100644 --- a/packages/ai-config/src/schema.ts +++ b/packages/ai-config/src/schema.ts @@ -149,6 +149,17 @@ export const snowflakeConfigSchema = z }) .strict(); +/** + * Databricks workspace host (NOT a chat base URL — the bridge derives the + * serving-endpoints / AI Gateway URL from it). Kept out of `baseUrl` so the + * per-model endpoint resolution never routes chat to the bare host. + */ +export const databricksConfigSchema = z + .object({ + host: z.string().optional(), + }) + .strict(); + /** Per-protocol base-URL overrides (partial — only specified protocols). */ export const endpointsSchema = z.record(protocolSchema, z.string().optional()); @@ -180,6 +191,7 @@ const CONNECTION_SECTION_SCHEMAS = { aws: awsConfigSchema, googleCloud: googleCloudConfigSchema, snowflake: snowflakeConfigSchema, + databricks: databricksConfigSchema, positaiLogin: positaiLoginConfigSchema, } as const; @@ -196,6 +208,7 @@ const allConnectionFields = { aws: awsConfigSchema.optional(), googleCloud: googleCloudConfigSchema.optional(), snowflake: snowflakeConfigSchema.optional(), + databricks: databricksConfigSchema.optional(), positaiLogin: positaiLoginConfigSchema.optional(), }; @@ -253,6 +266,7 @@ const BUILTIN_CONNECTION_SECTIONS = { "snowflake-cortex": ["snowflake"], "ms-foundry": [], deepseek: [], + databricks: ["databricks"], } as const satisfies Record; /** diff --git a/packages/ai-config/src/types.ts b/packages/ai-config/src/types.ts index b8e7d8f..f6cd144 100644 --- a/packages/ai-config/src/types.ts +++ b/packages/ai-config/src/types.ts @@ -126,6 +126,7 @@ export interface ResolvedConnection { aws?: { region?: string; profile?: string }; googleCloud?: { project?: string; location?: string }; snowflake?: { account?: string; host?: string; home?: string }; + databricks?: { host?: string }; } /** diff --git a/packages/ai-config/src/vocabulary.ts b/packages/ai-config/src/vocabulary.ts index f8ec8ef..13f3da5 100644 --- a/packages/ai-config/src/vocabulary.ts +++ b/packages/ai-config/src/vocabulary.ts @@ -33,6 +33,7 @@ export const BUILTIN_PROVIDER_IDS = [ "snowflake-cortex", "ms-foundry", "deepseek", + "databricks", ] as const; export type BuiltinProviderId = (typeof BUILTIN_PROVIDER_IDS)[number]; @@ -86,6 +87,7 @@ export const CLIENT_KIND_VALUES = [ "positai", "copilot", "ms-foundry", + "databricks", ] as const; export type ClientKind = (typeof CLIENT_KIND_VALUES)[number]; diff --git a/packages/ai-credentials/src/Backend.ts b/packages/ai-credentials/src/Backend.ts index bdfb8c5..10b962d 100644 --- a/packages/ai-credentials/src/Backend.ts +++ b/packages/ai-credentials/src/Backend.ts @@ -20,7 +20,12 @@ * providers resolve through {@link Backend.getCredentials}. */ -import type { Disposable } from "./CredentialProvider"; +import type { + CredentialMutation, + CredentialSourceInput, + CredentialStatus, + Disposable, +} from "./CredentialProvider"; import type { ProviderCredentials, TokenData } from "./types"; /** @@ -39,6 +44,70 @@ export interface OAuthProviderConfig { clientId: string; } +export interface AuthorizationCodeCallback { + code?: string; + error?: string; + errorDescription?: string; +} + +export interface PreparedAuthorizationCodeReceiver { + redirectUri: string; + waitForCallback(): Promise; + dispose(): void; +} + +/** Host adapter for a loopback callback (or an in-memory test receiver). */ +export interface AuthorizationCodeReceiver { + prepare(input: { + attemptId: string; + state: string; + timeoutMs: number; + }): Promise; +} + +/** Provider-neutral OAuth grant configuration consumed by the acquisition engine. */ +export type OAuthGrantConfig = + | { + grantType: "device-code"; + credentialBaseUrl?: string; + clientId: string; + scope: string; + deviceAuthorizationEndpoint: string; + tokenEndpoint: string; + } + | { + grantType: "authorization-code"; + credentialBaseUrl?: string; + clientId: string; + scope: string; + authorizationEndpoint: string; + tokenEndpoint: string; + receiver: AuthorizationCodeReceiver; + timeoutMs?: number; + challengeExpiresIn?: number; + } + | { + grantType: "client-credentials"; + credentialBaseUrl?: string; + clientId: string; + clientSecret: string; + scope?: string; + tokenEndpoint: string; + /** Non-secret identity used for process-memory token caching. */ + cacheKey: string; + }; + +export type CredentialSourceContext = + | { type: "oauth-device"; origin: "stored" | "implicit" } + | { type: "oauth-u2m"; origin: "stored"; workspaceHost: string } + | { + type: "oauth-m2m"; + origin: "stored" | "environment"; + workspaceHost: string; + clientId: string; + clientSecret: string; + }; + /** * Currently-stored OAuth tokens for a provider, as read back by the backend. * @@ -98,6 +167,34 @@ export interface OAuthBackendHooks { notifyReady(providerId: string): void; } +export type AuthenticationCommitResult = "committed" | "superseded"; + +/** Durable hooks used by the generalized acquisition engine. */ +export interface AcquisitionBackendHooks { + configForProvider(providerId: string): Promise; + readTokens(providerId: string): Promise; + beginAuthentication(providerId: string): Promise; + commitAuthentication( + providerId: string, + generation: string, + tokens: TokenData, + ): Promise; + finishAuthentication( + providerId: string, + generation: string, + error: string, + ): Promise; + persistRefreshedTokens(providerId: string, tokens: TokenData): Promise; + persistRefreshError(providerId: string, error: string): Promise; + withRefreshTransaction(providerId: string, operation: () => Promise): Promise; + shapeToken( + providerId: string, + accessToken: string, + config: OAuthGrantConfig, + ): ProviderCredentials; + notifyReady(providerId: string): void; +} + /** * The host-selected credential material seam. See the file header. */ @@ -115,4 +212,14 @@ export interface Backend { /** OAuth device-flow hooks; absent when the backend has no device flow. */ oauth?: OAuthBackendHooks; + + /** Generalized OAuth acquisition hooks used by upgraded store backends. */ + acquisition?: AcquisitionBackendHooks; +} + +export interface MutableBackend extends Backend { + mutateCredentials(providerId: string, mutation: CredentialMutation): Promise; + getCredentialStatus(providerId: string): Promise; + /** Resolve the active source without exposing its disk representation. */ + getCredentialSource(providerId: string): Promise; } diff --git a/packages/ai-credentials/src/CredentialProvider.ts b/packages/ai-credentials/src/CredentialProvider.ts index 1869051..7240d2b 100644 --- a/packages/ai-credentials/src/CredentialProvider.ts +++ b/packages/ai-credentials/src/CredentialProvider.ts @@ -4,6 +4,64 @@ import type { DeviceAuthInfo, ProviderCredentials } from "./types"; +export type AuthenticationChallenge = + | { + kind: "device-code"; + attemptId: string; + verificationUri: string; + verificationUriComplete: string; + userCode: string; + expiresIn: number; + } + | { + kind: "authorization-code"; + attemptId: string; + authorizationUrl: string; + expiresIn: number; + }; + +export type AuthenticationStartResult = + | { status: "started"; challenge: AuthenticationChallenge } + | { status: "already-in-progress" }; + +/** Strict semantic inputs accepted by the store-backed credential controller. */ +export type CredentialSourceInput = + | { type: "api-key"; apiKey: string; baseUrl?: string } + | { type: "oauth-device" } + | { type: "oauth-u2m"; workspaceHost: string } + | { + type: "oauth-m2m"; + clientId: string; + clientSecret: string; + workspaceHost: string; + } + | { type: "local"; endpoint: string } + | { + type: "aws-credentials"; + region: string; + profile?: string; + accessKeyId?: string; + secretAccessKey?: string; + sessionToken?: string; + } + | { type: "google-cloud"; project: string; location: string }; + +export type CredentialMutation = + | { kind: "replace"; source: CredentialSourceInput } + | { kind: "clear" }; + +export interface CredentialStatus { + configured: boolean; + authenticated: boolean; + readiness: "pending" | "ready" | "unauthenticated"; + source?: CredentialSourceInput["type"]; + origin?: "stored" | "environment"; + expiresAt?: string; + scope?: string; + error?: string; + metadata?: Record; +} + /** * Disposable handle returned by event subscriptions. */ @@ -33,6 +91,12 @@ export interface CredentialProvider { */ getCredentials(providerId: string): Promise; + /** Start an interactive authentication flow, if the provider supports one. */ + startAuthentication(providerId: string): Promise; + + /** Cancel one opaque authentication attempt. Unknown ids are ignored. */ + cancelAuthentication(attemptId: string): void; + /** * Current OAuth access token for a provider, refreshing if it is expired or * within the jittered refresh window. Null if the provider is not a @@ -50,3 +114,9 @@ export interface CredentialProvider { /** Subscribe to credential changes for the given provider ids. */ onDidChangeCredentials(callback: (providerIds: string[]) => void): Disposable; } + +/** Store-backed extension. Host-owned backends intentionally omit mutation. */ +export interface MutableCredentialProvider extends CredentialProvider { + mutateCredentials(providerId: string, mutation: CredentialMutation): Promise; + getCredentialStatus(providerId: string): Promise; +} diff --git a/packages/ai-credentials/src/__tests__/acquisition.test.ts b/packages/ai-credentials/src/__tests__/acquisition.test.ts new file mode 100644 index 0000000..1b1bd18 --- /dev/null +++ b/packages/ai-credentials/src/__tests__/acquisition.test.ts @@ -0,0 +1,411 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (C) 2026 Posit Software, PBC. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +import { mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; + +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +import type { + AuthorizationCodeCallback, + AuthorizationCodeReceiver, + CredentialSourceContext, + OAuthGrantConfig, + OAuthProviderConfig, + PreparedAuthorizationCodeReceiver, +} from "../Backend"; +import { createCredentialProvider } from "../createCredentialProvider"; +import { SingleFileStore } from "../store"; +import { createStoreBackend } from "../store-backend/StoreBackend"; +import type { StoredProviderCredentials } from "../store-backend/StoredProviderCredentials"; + +class TestReceiver implements AuthorizationCodeReceiver { + private resolveCallback?: (callback: AuthorizationCodeCallback) => void; + + prepare(): Promise { + return Promise.resolve({ + redirectUri: "http://127.0.0.1:8020/", + waitForCallback: () => + new Promise((resolve) => { + this.resolveCallback = resolve; + }), + dispose() {}, + }); + } + + complete(callback: AuthorizationCodeCallback): void { + this.resolveCallback?.(callback); + } +} + +const ok = (body: unknown): Response => + new Response(JSON.stringify(body), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + +describe("generalized store-backed acquisition", () => { + let directory: string; + let store: SingleFileStore; + let receiver: TestReceiver; + let generations: number; + + beforeEach(() => { + directory = mkdtempSync(join(tmpdir(), "acquisition-")); + store = new SingleFileStore({ filePath: join(directory, "data.json") }); + receiver = new TestReceiver(); + generations = 0; + }); + + afterEach(() => { + vi.unstubAllGlobals(); + rmSync(directory, { recursive: true, force: true }); + }); + + function createProvider( + env: Record = {}, + authorizationReceiver: AuthorizationCodeReceiver = receiver, + ) { + const backend = createStoreBackend({ + store, + env, + generationFactory: () => `generation-${++generations}`, + resolveAuthMethod: (providerId) => { + if (providerId === "databricks") return { authMethodId: "apikey" }; + if (providerId === "positai") return { authMethodId: "oauth" }; + return undefined; + }, + oauthConfigForProvider: ( + providerId: string, + source?: CredentialSourceContext, + ): OAuthGrantConfig | OAuthProviderConfig | undefined => { + if (providerId === "positai") { + return { authHost: "auth.test", clientId: "posit-ai", scope: "prism" }; + } + if (source?.type === "oauth-u2m") { + return { + grantType: "authorization-code", + clientId: "client", + scope: "all-apis offline_access", + authorizationEndpoint: `${source.workspaceHost}/authorize`, + tokenEndpoint: `${source.workspaceHost}/token`, + credentialBaseUrl: source.workspaceHost, + receiver: authorizationReceiver, + }; + } + if (source?.type === "oauth-m2m") { + return { + grantType: "client-credentials", + clientId: source.clientId, + clientSecret: source.clientSecret, + tokenEndpoint: `${source.workspaceHost}/token`, + credentialBaseUrl: source.workspaceHost, + cacheKey: `${source.workspaceHost}:${source.clientId}`, + }; + } + return undefined; + }, + }); + return createCredentialProvider({ backend }); + } + + it("completes authorization-code PKCE and rejects a genuinely concurrent local start", async () => { + const provider = createProvider(); + await provider.mutateCredentials("databricks", { + kind: "replace", + source: { type: "oauth-u2m", workspaceHost: "https://workspace.test" }, + }); + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue( + ok({ + access_token: "access", + refresh_token: "refresh", + expires_in: 3600, + token_type: "Bearer", + scope: "all-apis offline_access", + }), + ), + ); + + const [started, concurrent] = await Promise.all([ + provider.startAuthentication("databricks"), + provider.startAuthentication("databricks"), + ]); + expect(started.status).toBe("started"); + expect(concurrent).toEqual({ + status: "already-in-progress", + }); + const pending = await store.get("auth:databricks:apikey"); + expect(pending).toMatchObject({ readiness: "pending", authenticated: false }); + expect(pending?.oauthAuth?.tokenData).toBeUndefined(); + + receiver.complete({ code: "code" }); + await vi.waitFor(async () => { + expect(await provider.getCredentials("databricks")).toEqual({ + type: "apikey", + apiKey: "access", + baseUrl: "https://workspace.test", + }); + }); + }); + + it("does not let a stale callback resurrect credentials after clear", async () => { + const provider = createProvider(); + await provider.mutateCredentials("databricks", { + kind: "replace", + source: { type: "oauth-u2m", workspaceHost: "https://workspace.test" }, + }); + vi.stubGlobal( + "fetch", + vi + .fn() + .mockResolvedValue( + ok({ access_token: "stale", refresh_token: "stale-refresh", expires_in: 3600 }), + ), + ); + await provider.startAuthentication("databricks"); + await provider.mutateCredentials("databricks", { kind: "clear" }); + receiver.complete({ code: "late" }); + await new Promise((resolve) => setTimeout(resolve, 0)); + + expect(await provider.getCredentials("databricks")).toBeNull(); + expect(await store.get("auth:databricks:apikey")).toMatchObject({ + readiness: "unauthenticated", + configured: false, + }); + }); + + it("cancels by opaque attempt ID and replaces pending state with a fresh terminal generation", async () => { + const provider = createProvider(); + await provider.mutateCredentials("databricks", { + kind: "replace", + source: { type: "oauth-u2m", workspaceHost: "https://workspace.test" }, + }); + const started = await provider.startAuthentication("databricks"); + if (started.status !== "started") throw new Error("Expected authentication to start"); + const pending = await store.get("auth:databricks:apikey"); + provider.cancelAuthentication(started.challenge.attemptId); + + await vi.waitFor(async () => { + const terminal = await store.get("auth:databricks:apikey"); + expect(terminal).toMatchObject({ + readiness: "unauthenticated", + authenticated: false, + error: "cancelled", + }); + expect(terminal?.generation).not.toBe(pending?.generation); + }); + receiver.complete({ code: "late" }); + expect(await provider.getCredentials("databricks")).toBeNull(); + }); + + it("renews environment M2M in memory without persisting secrets or tokens", async () => { + const provider = createProvider({ + DATABRICKS_CLIENT_ID: "client", + DATABRICKS_CLIENT_SECRET: "secret", + DATABRICKS_HOST: "https://workspace.test", + }); + vi.stubGlobal( + "fetch", + vi + .fn() + .mockResolvedValue( + ok({ access_token: "m2m-access", expires_in: 3600, token_type: "Bearer" }), + ), + ); + + expect(await provider.getCredentials("databricks")).toEqual({ + type: "apikey", + apiKey: "m2m-access", + baseUrl: "https://workspace.test", + }); + expect(await store.keys()).toEqual([]); + expect(await provider.getCredentialStatus("databricks")).toMatchObject({ + source: "oauth-m2m", + origin: "environment", + }); + }); + + it("treats a legacy generationless PAT as an explicit stored source", async () => { + await store.set("auth:databricks:apikey", { + apiKeyAuth: { apiKey: "legacy", baseUrl: "https://legacy.test" }, + }); + const provider = createProvider({ + DATABRICKS_CLIENT_ID: "client", + DATABRICKS_CLIENT_SECRET: "secret", + DATABRICKS_HOST: "https://environment.test", + }); + expect(await provider.getCredentials("databricks")).toEqual({ + type: "apikey", + apiKey: "legacy", + baseUrl: "https://legacy.test", + }); + }); + + it("rejects a stale process across clear, a generationless legacy write, and a later attempt", async () => { + const firstReceiver = new TestReceiver(); + const secondReceiver = new TestReceiver(); + const firstProcess = createProvider({}, firstReceiver); + const secondProcess = createProvider({}, secondReceiver); + await firstProcess.mutateCredentials("databricks", { + kind: "replace", + source: { type: "oauth-u2m", workspaceHost: "https://workspace.test" }, + }); + await firstProcess.startAuthentication("databricks"); + await secondProcess.mutateCredentials("databricks", { kind: "clear" }); + await store.set("auth:databricks:apikey", { + apiKeyAuth: { apiKey: "legacy", baseUrl: "https://workspace.test" }, + }); + await secondProcess.mutateCredentials("databricks", { + kind: "replace", + source: { type: "oauth-u2m", workspaceHost: "https://workspace.test" }, + }); + await secondProcess.startAuthentication("databricks"); + vi.stubGlobal( + "fetch", + vi.fn(async (_input: string | URL | Request, init?: RequestInit) => { + const body = new URLSearchParams(typeof init?.body === "string" ? init.body : ""); + const code = body.get("code"); + return ok({ + access_token: code === "first" ? "stale" : "current", + refresh_token: `${code}-refresh`, + expires_in: 3600, + }); + }), + ); + + firstReceiver.complete({ code: "first" }); + secondReceiver.complete({ code: "second" }); + await vi.waitFor(async () => { + expect(await secondProcess.getCredentials("databricks")).toMatchObject({ + apiKey: "current", + }); + }); + expect(await firstProcess.getCredentials("databricks")).toMatchObject({ apiKey: "current" }); + }); + + describe("Posit AI device authentication through the store backend", () => { + beforeEach(() => vi.useFakeTimers()); + afterEach(() => vi.useRealTimers()); + + function deviceCodeResponse(): Response { + return ok({ + user_code: "WXYZ", + verification_uri: "https://auth.test/device", + verification_uri_complete: "https://auth.test/device?code=WXYZ", + device_code: "device-code", + interval: 1, + expires_in: 900, + }); + } + + it("commits success and records cancellation and errors as terminal generations", async () => { + const provider = createProvider(); + const fetchMock = vi + .fn() + .mockResolvedValueOnce(deviceCodeResponse()) + .mockResolvedValueOnce( + ok({ + access_token: "posit-access", + refresh_token: "posit-refresh", + expires_in: 3600, + token_type: "Bearer", + scope: "prism", + }), + ); + vi.stubGlobal("fetch", fetchMock); + + const successful = await provider.startAuthentication("positai"); + expect(successful.status).toBe("started"); + await vi.advanceTimersByTimeAsync(1000); + await vi.waitFor(async () => { + expect(await provider.getCredentials("positai")).toEqual({ + type: "oauth", + accessToken: "posit-access", + }); + }); + + fetchMock.mockResolvedValueOnce(deviceCodeResponse()); + const cancelled = await provider.startAuthentication("positai"); + if (cancelled.status !== "started") throw new Error("Expected authentication to start"); + provider.cancelAuthentication(cancelled.challenge.attemptId); + await vi.waitFor(async () => { + expect(await store.get("auth:positai:oauth")).toMatchObject({ + readiness: "unauthenticated", + error: "cancelled", + }); + }); + + fetchMock.mockResolvedValueOnce(deviceCodeResponse()).mockResolvedValueOnce( + new Response(JSON.stringify({ error: "access_denied" }), { + status: 400, + headers: { "Content-Type": "application/json" }, + }), + ); + await provider.startAuthentication("positai"); + await vi.advanceTimersByTimeAsync(1000); + await vi.waitFor(async () => { + expect(await store.get("auth:positai:oauth")).toMatchObject({ + readiness: "unauthenticated", + error: "access_denied", + }); + }); + }); + + it("shares one attempt across generic and compatibility surfaces", async () => { + const provider = createProvider(); + vi.stubGlobal("fetch", vi.fn().mockResolvedValueOnce(deviceCodeResponse())); + + await provider.startDeviceAuth("positai"); + expect(await provider.startAuthentication("positai")).toEqual({ + status: "already-in-progress", + }); + provider.cancelDeviceAuth("positai"); + await provider.dispose(); + }); + + it("does not let compatibility polling resurrect credentials after clear", async () => { + const provider = createProvider(); + const fetchMock = vi + .fn() + .mockResolvedValueOnce(deviceCodeResponse()) + .mockResolvedValueOnce( + ok({ + access_token: "stale", + refresh_token: "stale-refresh", + expires_in: 3600, + token_type: "Bearer", + scope: "prism", + }), + ); + vi.stubGlobal("fetch", fetchMock); + + await provider.startDeviceAuth("positai"); + await provider.mutateCredentials("positai", { kind: "clear" }); + await vi.advanceTimersByTimeAsync(5000); + expect(fetchMock).toHaveBeenCalledOnce(); + expect(await provider.getCredentials("positai")).toBeNull(); + expect(await store.get("auth:positai:oauth")).toMatchObject({ + readiness: "unauthenticated", + configured: false, + }); + await provider.dispose(); + }); + + it("durably terminates pending authentication during graceful disposal", async () => { + const provider = createProvider(); + vi.stubGlobal("fetch", vi.fn().mockResolvedValueOnce(deviceCodeResponse())); + await provider.startAuthentication("positai"); + + await provider.dispose(); + + expect(await store.get("auth:positai:oauth")).toMatchObject({ + readiness: "unauthenticated", + error: "cancelled", + }); + }); + }); +}); diff --git a/packages/ai-credentials/src/__tests__/databricks-oauth.test.ts b/packages/ai-credentials/src/__tests__/databricks-oauth.test.ts new file mode 100644 index 0000000..4be4573 --- /dev/null +++ b/packages/ai-credentials/src/__tests__/databricks-oauth.test.ts @@ -0,0 +1,32 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (C) 2026 Posit Software, PBC. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +import { describe, expect, it } from "vitest"; + +import { normalizeDatabricksWorkspaceHost } from "../databricks-oauth"; +import { normalizeDatabricksHost } from "../types"; + +describe("Databricks OAuth workspace validation", () => { + it("normalizes secure workspace hosts", () => { + expect(normalizeDatabricksWorkspaceHost("workspace.example.com/path?q=1")).toBe( + "https://workspace.example.com", + ); + }); + + it.each(["http://workspace.example.com", "http://localhost:8080", "http://127.0.0.1:8080"])( + "rejects insecure production workspace URL %s", + (workspaceHost) => { + expect(() => normalizeDatabricksWorkspaceHost(workspaceHost)).toThrow( + "Databricks workspace URL must use HTTPS", + ); + }, + ); + + it("uses the same HTTPS-only parser from the browser-safe types entrypoint", () => { + expect(normalizeDatabricksHost("workspace.example.com/path?q=1")).toBe( + "https://workspace.example.com", + ); + expect(() => normalizeDatabricksHost("http://workspace.example.com")).toThrow("must use HTTPS"); + }); +}); diff --git a/packages/ai-credentials/src/acquisition.ts b/packages/ai-credentials/src/acquisition.ts new file mode 100644 index 0000000..f04851a --- /dev/null +++ b/packages/ai-credentials/src/acquisition.ts @@ -0,0 +1,599 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (C) 2026 Posit Software, PBC. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +import type { + AcquisitionBackendHooks, + OAuthGrantConfig, + PreparedAuthorizationCodeReceiver, + StoredOAuthTokens, +} from "./Backend"; +import type { AuthenticationStartResult } from "./CredentialProvider"; +import type { DeviceAuthInfo, Logger, ProviderCredentials, TokenData } from "./types"; + +interface TokenResponse { + access_token?: unknown; + refresh_token?: unknown; + expires_in?: unknown; + token_type?: unknown; + scope?: unknown; +} + +interface ActiveAttempt { + attemptId: string; + providerId: string; + generation: string; + controller: AbortController; + receiver?: PreparedAuthorizationCodeReceiver; + terminalPromise?: Promise; +} + +interface DeviceAuthenticationStart { + result: AuthenticationStartResult; + info: DeviceAuthInfo; +} + +const DEFAULT_AUTHORIZATION_TIMEOUT_MS = 5 * 60 * 1000; + +/** Provider-neutral device/code/client-credentials acquisition state machine. */ +export class AcquisitionEngine { + private readonly activeByProvider = new Map(); + private readonly activeById = new Map(); + private readonly startingProviders = new Set(); + private readonly refreshPromises = new Map>(); + private readonly clientCredentialTokens = new Map(); + private readonly refreshJitterMinutes = 4 + Math.random() * 2; + private readonly startPromises = new Set>(); + private readonly terminalPromises = new Set>(); + private disposed = false; + private disposalPromise: Promise | undefined; + + constructor( + private readonly hooks: AcquisitionBackendHooks, + private readonly logger?: Logger, + ) {} + + async getCredentials( + providerId: string, + ): Promise<{ handled: boolean; credentials: ProviderCredentials | null }> { + const config = await this.hooks.configForProvider(providerId); + if (!config) return { handled: false, credentials: null }; + + if (config.grantType === "client-credentials") { + return { handled: true, credentials: await this.getClientCredentials(providerId, config) }; + } + + const tokens = await this.hooks.readTokens(providerId); + if (!tokens) return { handled: true, credentials: null }; + if (!this.isExpiring(tokens)) { + return { + handled: true, + credentials: this.hooks.shapeToken(providerId, tokens.accessToken, config), + }; + } + + return { handled: true, credentials: await this.refreshStored(providerId, config) }; + } + + startAuthentication(providerId: string): Promise { + if (this.disposed) return Promise.reject(new Error("Credential provider is disposed")); + if (this.activeByProvider.has(providerId) || this.startingProviders.has(providerId)) { + return Promise.resolve({ status: "already-in-progress" }); + } + this.startingProviders.add(providerId); + const start = this.startReserved(providerId, false).then((value) => value.result); + return this.trackStart(providerId, start); + } + + startDeviceAuthentication(providerId: string): Promise { + if (this.disposed) return Promise.reject(new Error("Credential provider is disposed")); + if (this.activeByProvider.has(providerId) || this.startingProviders.has(providerId)) { + return Promise.reject(new Error(`Authentication is already in progress for ${providerId}`)); + } + this.startingProviders.add(providerId); + const start = this.startReserved(providerId, true).then((value) => { + if (!value.device) { + throw new Error(`OAuth device auth not supported for provider: ${providerId}`); + } + return value.device.info; + }); + return this.trackStart(providerId, start); + } + + private async startReserved( + providerId: string, + deviceOnly: boolean, + ): Promise<{ result: AuthenticationStartResult; device?: DeviceAuthenticationStart }> { + let attempt: ActiveAttempt | undefined; + try { + const config = await this.hooks.configForProvider(providerId); + if (this.disposed) throw new Error("Credential provider is disposed"); + if (!config || config.grantType === "client-credentials") { + throw new Error(`Interactive authentication is not supported for provider: ${providerId}`); + } + if (deviceOnly && config.grantType !== "device-code") { + throw new Error(`OAuth device auth not supported for provider: ${providerId}`); + } + + const attemptId = randomOpaque(16); + const generation = await this.hooks.beginAuthentication(providerId); + if (this.disposed) { + await this.hooks.finishAuthentication(providerId, generation, "cancelled"); + throw new Error("Credential provider is disposed"); + } + attempt = { + attemptId, + providerId, + generation, + controller: new AbortController(), + }; + this.activeByProvider.set(providerId, attempt); + this.activeById.set(attemptId, attempt); + + if (config.grantType === "device-code") { + const device = await this.startDeviceCode(attempt, config); + return { result: device.result, device }; + } + return { result: await this.startAuthorizationCode(attempt, config) }; + } catch (error) { + if (attempt) { + await this.terminateAttempt(attempt, errorCode(error)); + } + throw error; + } + } + + cancelAuthentication(attemptId: string): void { + const attempt = this.activeById.get(attemptId); + if (!attempt) return; + void this.terminateAttempt(attempt, "cancelled"); + } + + cancelProvider(providerId: string, persistTerminal = true): void { + const attempt = this.activeByProvider.get(providerId); + if (!attempt) return; + if (persistTerminal) { + void this.terminateAttempt(attempt, "cancelled"); + } else { + attempt.controller.abort(); + attempt.receiver?.dispose(); + this.removeAttempt(attempt); + } + } + + dispose(): Promise { + if (this.disposalPromise) return this.disposalPromise; + this.disposed = true; + const terminalWrites = [...this.activeById.values()].map((attempt) => + this.terminateAttempt(attempt, "cancelled"), + ); + this.disposalPromise = Promise.allSettled([ + ...this.startPromises, + ...this.terminalPromises, + ...terminalWrites, + ]).then(() => undefined); + return this.disposalPromise; + } + + private async startDeviceCode( + attempt: ActiveAttempt, + config: Extract, + ): Promise { + const response = await postForm( + config.deviceAuthorizationEndpoint, + { + scope: config.scope, + client_id: config.clientId, + }, + attempt.controller.signal, + ); + const data = await readObject(response, "Device authorization"); + const userCode = requiredString(data, "user_code"); + const verificationUri = requiredString(data, "verification_uri"); + const verificationUriComplete = requiredString(data, "verification_uri_complete"); + const deviceCode = requiredString(data, "device_code"); + const interval = requiredPositiveNumber(data, "interval"); + const expiresIn = requiredPositiveNumber(data, "expires_in"); + + void this.pollDeviceCode(attempt, config, deviceCode, interval).catch((error: unknown) => { + this.logger?.error( + `[ai-credentials] device authentication failed for ${attempt.providerId}`, + error, + ); + }); + + return { + result: { + status: "started", + challenge: { + kind: "device-code", + attemptId: attempt.attemptId, + verificationUri, + verificationUriComplete, + userCode, + expiresIn, + }, + }, + info: { + verificationUri, + verificationUriComplete, + userCode, + deviceCode, + interval, + expiresIn, + }, + }; + } + + private async startAuthorizationCode( + attempt: ActiveAttempt, + config: Extract, + ): Promise { + const state = randomOpaque(32); + const verifier = randomOpaque(64); + const challenge = await sha256Base64Url(verifier); + const timeoutMs = config.timeoutMs ?? DEFAULT_AUTHORIZATION_TIMEOUT_MS; + const receiver = await config.receiver.prepare({ + attemptId: attempt.attemptId, + state, + timeoutMs, + }); + attempt.receiver = receiver; + + const url = new URL(config.authorizationEndpoint); + url.search = new URLSearchParams({ + client_id: config.clientId, + redirect_uri: receiver.redirectUri, + response_type: "code", + scope: config.scope, + state, + code_challenge: challenge, + code_challenge_method: "S256", + }).toString(); + + void this.completeAuthorizationCode(attempt, config, verifier).catch((error: unknown) => { + this.logger?.error( + `[ai-credentials] authorization-code authentication failed for ${attempt.providerId}`, + error, + ); + }); + + return { + status: "started", + challenge: { + kind: "authorization-code", + attemptId: attempt.attemptId, + authorizationUrl: url.toString(), + expiresIn: Math.floor(config.challengeExpiresIn ?? timeoutMs / 1000), + }, + }; + } + + private async completeAuthorizationCode( + attempt: ActiveAttempt, + config: Extract, + verifier: string, + ): Promise { + try { + const receiver = attempt.receiver; + if (!receiver) throw new Error("authorization_callback_missing"); + const callback = await receiver.waitForCallback(); + if (callback.error) { + throw new Error(callback.errorDescription || callback.error); + } + if (!callback.code) throw new Error("authorization_code_missing"); + if (!this.isCurrent(attempt)) return; + + const response = await postForm( + config.tokenEndpoint, + { + grant_type: "authorization_code", + client_id: config.clientId, + code: callback.code, + code_verifier: verifier, + redirect_uri: receiver.redirectUri, + }, + attempt.controller.signal, + ); + const tokens = await tokenData(response, true); + const committed = await this.hooks.commitAuthentication( + attempt.providerId, + attempt.generation, + tokens, + ); + if (committed === "committed") this.hooks.notifyReady(attempt.providerId); + } catch (error) { + if (this.isCurrent(attempt)) { + await this.hooks.finishAuthentication( + attempt.providerId, + attempt.generation, + errorCode(error), + ); + } + } finally { + attempt.receiver?.dispose(); + this.removeAttempt(attempt); + } + } + + private async pollDeviceCode( + attempt: ActiveAttempt, + config: Extract, + deviceCode: string, + intervalSeconds: number, + ): Promise { + let intervalMs = intervalSeconds * 1000; + try { + while (this.isCurrent(attempt) && !attempt.controller.signal.aborted) { + await sleep(intervalMs, attempt.controller.signal); + const response = await postForm( + config.tokenEndpoint, + { + grant_type: "urn:ietf:params:oauth:grant-type:device_code", + client_id: config.clientId, + scope: config.scope, + device_code: deviceCode, + }, + attempt.controller.signal, + true, + ); + if (response.ok) { + const tokens = await tokenData(response, true); + const committed = await this.hooks.commitAuthentication( + attempt.providerId, + attempt.generation, + tokens, + ); + if (committed === "committed") this.hooks.notifyReady(attempt.providerId); + return; + } + const body = await safeObject(response); + const code = typeof body.error === "string" ? body.error : `http_${response.status}`; + if (code === "authorization_pending") continue; + if (code === "slow_down") { + intervalMs += 5000; + continue; + } + throw new Error(code); + } + } catch (error) { + if (this.isCurrent(attempt) && !attempt.controller.signal.aborted) { + await this.hooks.finishAuthentication( + attempt.providerId, + attempt.generation, + errorCode(error), + ); + } + } finally { + this.removeAttempt(attempt); + } + } + + private refreshStored( + providerId: string, + config: Exclude, + ): Promise { + const existing = this.refreshPromises.get(providerId); + if (existing) return existing; + const promise = this.hooks.withRefreshTransaction(providerId, async () => { + const current = await this.hooks.readTokens(providerId); + if (!current) return null; + if (!this.isExpiring(current, 2)) { + return this.hooks.shapeToken(providerId, current.accessToken, config); + } + try { + const response = await postForm(config.tokenEndpoint, { + grant_type: "refresh_token", + client_id: config.clientId, + refresh_token: current.refreshToken, + ...(config.scope ? { scope: config.scope } : {}), + }); + const refreshed = await tokenData(response, false, current.refreshToken); + await this.hooks.persistRefreshedTokens(providerId, refreshed); + this.hooks.notifyReady(providerId); + return this.hooks.shapeToken(providerId, refreshed.accessToken, config); + } catch (error) { + await this.hooks.persistRefreshError(providerId, "refresh_failed"); + this.logger?.error(`[ai-credentials] refresh failed for ${providerId}`, error); + return null; + } + }); + this.refreshPromises.set(providerId, promise); + return promise.finally(() => this.refreshPromises.delete(providerId)); + } + + private async getClientCredentials( + providerId: string, + config: Extract, + ): Promise { + const cached = this.clientCredentialTokens.get(config.cacheKey); + if (cached && !this.isExpiring(cached)) { + return this.hooks.shapeToken(providerId, cached.accessToken, config); + } + + const mutexKey = `${providerId}:${config.cacheKey}`; + const existing = this.refreshPromises.get(mutexKey); + if (existing) return existing; + const promise = (async () => { + try { + const response = await postForm(config.tokenEndpoint, { + grant_type: "client_credentials", + client_id: config.clientId, + client_secret: config.clientSecret, + ...(config.scope ? { scope: config.scope } : {}), + }); + const tokens = await tokenData(response, false, ""); + this.clientCredentialTokens.set(config.cacheKey, toStored(tokens)); + return this.hooks.shapeToken(providerId, tokens.accessToken, config); + } catch (error) { + this.logger?.error( + `[ai-credentials] client-credentials renewal failed for ${providerId}`, + error, + ); + return null; + } + })(); + this.refreshPromises.set(mutexKey, promise); + return promise.finally(() => this.refreshPromises.delete(mutexKey)); + } + + private isExpiring(tokens: StoredOAuthTokens, fixedMinutes?: number): boolean { + const minutes = fixedMinutes ?? this.refreshJitterMinutes; + return new Date(tokens.expiresAt).getTime() <= Date.now() + minutes * 60 * 1000; + } + + private isCurrent(attempt: ActiveAttempt): boolean { + return this.activeById.get(attempt.attemptId) === attempt; + } + + private trackStart(providerId: string, start: Promise): Promise { + this.startPromises.add(start); + return start.finally(() => { + this.startPromises.delete(start); + this.startingProviders.delete(providerId); + }); + } + + private terminateAttempt(attempt: ActiveAttempt, error: string): Promise { + attempt.controller.abort(); + attempt.receiver?.dispose(); + this.removeAttempt(attempt); + if (!attempt.terminalPromise) { + const terminalPromise = this.hooks + .finishAuthentication(attempt.providerId, attempt.generation, error) + .then(() => undefined); + attempt.terminalPromise = terminalPromise; + this.terminalPromises.add(terminalPromise); + void terminalPromise.then( + () => this.terminalPromises.delete(terminalPromise), + () => this.terminalPromises.delete(terminalPromise), + ); + } + return attempt.terminalPromise; + } + + private removeAttempt(attempt: ActiveAttempt): void { + if (this.activeById.get(attempt.attemptId) === attempt) { + this.activeById.delete(attempt.attemptId); + } + if (this.activeByProvider.get(attempt.providerId) === attempt) { + this.activeByProvider.delete(attempt.providerId); + } + } +} + +async function postForm( + url: string, + params: Record, + signal?: AbortSignal, + allowError = false, +): Promise { + const response = await fetch(url, { + method: "POST", + headers: { "Content-Type": "application/x-www-form-urlencoded" }, + body: new URLSearchParams(params).toString(), + signal, + }); + if (!allowError && !response.ok) { + throw new Error(`oauth_http_${response.status}`); + } + return response; +} + +async function tokenData( + response: Response, + requireRefreshToken: boolean, + refreshTokenFallback?: string, +): Promise { + const body = await readObject(response, "Token exchange"); + const accessToken = requiredString(body, "access_token"); + const refreshTokenValue = + typeof body.refresh_token === "string" ? body.refresh_token : refreshTokenFallback; + if (requireRefreshToken && !refreshTokenValue) throw new Error("malformed_refresh_token"); + return { + accessToken, + refreshToken: refreshTokenValue ?? "", + expiresIn: requiredPositiveNumber(body, "expires_in"), + tokenType: typeof body.token_type === "string" ? body.token_type : "Bearer", + scope: typeof body.scope === "string" ? body.scope : "", + }; +} + +function toStored(tokens: TokenData): StoredOAuthTokens { + return { + accessToken: tokens.accessToken, + refreshToken: tokens.refreshToken, + expiresAt: new Date(Date.now() + tokens.expiresIn * 1000).toISOString(), + tokenType: tokens.tokenType, + scope: tokens.scope, + }; +} + +async function readObject(response: Response, label: string): Promise> { + if (!response.ok) + throw new Error(`${label.toLowerCase().replaceAll(" ", "_")}_${response.status}`); + const value: unknown = await response.json(); + if (typeof value !== "object" || value === null) throw new Error("malformed_oauth_response"); + return value as Record; +} + +async function safeObject(response: Response): Promise> { + try { + const value: unknown = await response.json(); + return typeof value === "object" && value !== null ? (value as Record) : {}; + } catch { + return {}; + } +} + +function requiredString(value: Record, key: string): string { + const field = value[key]; + if (typeof field !== "string" || field.length === 0) throw new Error(`malformed_${key}`); + return field; +} + +function requiredPositiveNumber(value: Record, key: string): number { + const field = value[key]; + if (typeof field !== "number" || !Number.isFinite(field) || field <= 0) { + throw new Error(`malformed_${key}`); + } + return field; +} + +function randomOpaque(bytes: number): string { + const values = new Uint8Array(bytes); + globalThis.crypto.getRandomValues(values); + return [...values].map((value) => value.toString(16).padStart(2, "0")).join(""); +} + +async function sha256Base64Url(value: string): Promise { + const digest = await globalThis.crypto.subtle.digest("SHA-256", new TextEncoder().encode(value)); + const bytes = new Uint8Array(digest); + let binary = ""; + for (const byte of bytes) binary += String.fromCharCode(byte); + return btoa(binary).replaceAll("+", "-").replaceAll("/", "_").replace(/=+$/, ""); +} + +function errorCode(error: unknown): string { + if (error instanceof Error && error.name === "AbortError") return "cancelled"; + if (error instanceof Error && error.message) return error.message; + return "authentication_failed"; +} + +function sleep(ms: number, signal: AbortSignal): Promise { + return new Promise((resolve, reject) => { + if (signal.aborted) { + reject(new DOMException("Aborted", "AbortError")); + return; + } + const timer = setTimeout(resolve, ms); + signal.addEventListener( + "abort", + () => { + clearTimeout(timer); + reject(new DOMException("Aborted", "AbortError")); + }, + { once: true }, + ); + }); +} diff --git a/packages/ai-credentials/src/createCredentialProvider.ts b/packages/ai-credentials/src/createCredentialProvider.ts index fc33d03..729fe93 100644 --- a/packages/ai-credentials/src/createCredentialProvider.ts +++ b/packages/ai-credentials/src/createCredentialProvider.ts @@ -6,19 +6,27 @@ * createCredentialProvider — the deep resolver seam over a host-selected backend. * * Wraps any {@link Backend} into the full {@link CredentialProvider} resolver - * surface. The root owns the device-flow/refresh state machine ({@link OAuthEngine}); - * the backend supplies material, OAuth config, and token persistence. + * surface. The root owns the device-flow/refresh state machine; the backend + * supplies material, OAuth config, and token persistence. * * Routing: * - `getCredentials`: OAuth providers (backend exposes `oauth` config for the id) * route through `getAccessToken` and wrap the token; everything else defers to * `backend.getCredentials`. - * - `getAccessToken` / `startDeviceAuth`: delegate to the engine, which returns - * null / throws for providers with no device-flow config. + * - `getAccessToken` / `startDeviceAuth`: use the generalized acquisition + * controller when available, otherwise the legacy device-only engine. */ -import type { Backend } from "./Backend"; -import type { CredentialProvider, Disposable } from "./CredentialProvider"; +import { AcquisitionEngine } from "./acquisition"; +import type { Backend, MutableBackend } from "./Backend"; +import type { + AuthenticationStartResult, + CredentialMutation, + CredentialProvider, + CredentialStatus, + Disposable, + MutableCredentialProvider, +} from "./CredentialProvider"; import { OAuthEngine } from "./device-auth"; import type { DeviceAuthInfo, Logger, ProviderCredentials } from "./types"; @@ -40,22 +48,48 @@ export interface CredentialProviderHandle extends CredentialProvider { * No-op when the backend has no device flow or no polling is active. */ cancelDeviceAuth(providerId: string): void; - dispose(): void; + dispose(): Promise; } +export interface MutableCredentialProviderHandle + extends CredentialProviderHandle, MutableCredentialProvider {} + +export function createCredentialProvider( + options: CreateCredentialProviderOptions & { backend: MutableBackend }, +): MutableCredentialProviderHandle; +export function createCredentialProvider( + options: CreateCredentialProviderOptions, +): CredentialProviderHandle; + /** Build a {@link CredentialProvider} over the given backend. */ export function createCredentialProvider( options: CreateCredentialProviderOptions, ): CredentialProviderHandle { const { backend, logger } = options; - const engine = backend.oauth ? new OAuthEngine(backend.oauth, logger) : undefined; + const acquisition = backend.acquisition + ? new AcquisitionEngine(backend.acquisition, logger) + : undefined; + const legacyEngine = + !acquisition && backend.oauth ? new OAuthEngine(backend.oauth, logger) : undefined; async function getAccessToken(providerId: string): Promise { - if (!engine) return null; - return engine.getAccessToken(providerId); + if (acquisition) { + const result = await acquisition.getCredentials(providerId); + if (result.handled) { + if (result.credentials?.type === "oauth") return result.credentials.accessToken; + if (result.credentials?.type === "apikey") return result.credentials.apiKey; + return null; + } + } + if (!legacyEngine) return null; + return legacyEngine.getAccessToken(providerId); } async function getCredentials(providerId: string): Promise { + if (acquisition) { + const result = await acquisition.getCredentials(providerId); + if (result.handled) return result.credentials; + } // Device-flow OAuth providers (backend exposes config) resolve through the // engine so refresh is handled. Non-OAuth (and OAuth-via-vscode) defer to // the backend. @@ -66,13 +100,37 @@ export function createCredentialProvider( return backend.getCredentials(providerId); } + function startAuthentication(providerId: string): Promise { + if (acquisition) return acquisition.startAuthentication(providerId); + return startDeviceAuth(providerId).then((info) => ({ + status: "started" as const, + challenge: { + kind: "device-code" as const, + attemptId: providerId, + verificationUri: info.verificationUri, + verificationUriComplete: info.verificationUriComplete, + userCode: info.userCode, + expiresIn: info.expiresIn, + }, + })); + } + + function cancelAuthentication(attemptId: string): void { + if (acquisition) { + acquisition.cancelAuthentication(attemptId); + return; + } + cancelDeviceAuth(attemptId); + } + function startDeviceAuth(providerId: string): Promise { - if (!engine) { + if (acquisition) return acquisition.startDeviceAuthentication(providerId); + if (!legacyEngine) { return Promise.reject( new Error(`OAuth device auth not supported for provider: ${providerId}`), ); } - return engine.startDeviceAuth(providerId); + return legacyEngine.startDeviceAuth(providerId); } function onDidChangeCredentials(callback: (providerIds: string[]) => void): Disposable { @@ -80,19 +138,46 @@ export function createCredentialProvider( } function cancelDeviceAuth(providerId: string): void { - engine?.cancelPolling(providerId); + if (acquisition) { + acquisition.cancelProvider(providerId); + return; + } + legacyEngine?.cancelPolling(providerId); } - function dispose(): void { - engine?.dispose(); + async function dispose(): Promise { + legacyEngine?.dispose(); + await acquisition?.dispose(); } - return { + const result: CredentialProviderHandle = { getCredentials, + startAuthentication, + cancelAuthentication, getAccessToken, startDeviceAuth, onDidChangeCredentials, cancelDeviceAuth, dispose, }; + + if (isMutableBackend(backend)) { + const mutable = result as MutableCredentialProviderHandle; + mutable.mutateCredentials = async ( + providerId: string, + mutation: CredentialMutation, + ): Promise => { + acquisition?.cancelProvider(providerId, false); + await backend.mutateCredentials(providerId, mutation); + }; + mutable.getCredentialStatus = (providerId: string): Promise => + backend.getCredentialStatus(providerId); + return mutable; + } + + return result; +} + +function isMutableBackend(backend: Backend): backend is MutableBackend { + return "mutateCredentials" in backend && "getCredentialStatus" in backend; } diff --git a/packages/ai-credentials/src/databricks-oauth.ts b/packages/ai-credentials/src/databricks-oauth.ts new file mode 100644 index 0000000..756dad5 --- /dev/null +++ b/packages/ai-credentials/src/databricks-oauth.ts @@ -0,0 +1,122 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (C) 2026 Posit Software, PBC. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +import type { AuthorizationCodeReceiver, OAuthGrantConfig } from "./Backend"; +import { normalizeDatabricksHost } from "./types"; + +export interface DatabricksOidcEndpoints { + authorizationEndpoint: string; + tokenEndpoint: string; +} + +const endpointDiscovery = new Map>(); + +/** Normalize and validate a workspace URL without retaining path/query fragments. */ +export function normalizeDatabricksWorkspaceHost(raw: string): string { + return normalizeDatabricksHost(raw); +} + +/** Discover workspace OAuth endpoints, falling back to the documented workspace paths. */ +export async function discoverDatabricksOidcEndpoints( + workspaceHost: string, +): Promise { + const host = normalizeDatabricksWorkspaceHost(workspaceHost); + const existing = endpointDiscovery.get(host); + if (existing) return existing; + + const discovery = discoverEndpoints(host); + endpointDiscovery.set(host, discovery); + try { + return await discovery; + } catch (error) { + endpointDiscovery.delete(host); + throw error; + } +} + +async function discoverEndpoints(host: string): Promise { + const response = await fetch(`${host}/.well-known/openid-configuration`); + if (response.ok) { + const document: unknown = await response.json(); + if ( + hasStringProperty(document, "authorization_endpoint") && + hasStringProperty(document, "token_endpoint") + ) { + return { + authorizationEndpoint: validateEndpoint(host, document.authorization_endpoint), + tokenEndpoint: validateEndpoint(host, document.token_endpoint), + }; + } + } + return { + authorizationEndpoint: `${host}/oidc/v1/authorize`, + tokenEndpoint: `${host}/oidc/v1/token`, + }; +} + +export async function createDatabricksAuthorizationCodeGrant(input: { + workspaceHost: string; + receiver: AuthorizationCodeReceiver; + clientId: string; + scope: string; + timeoutMs?: number; +}): Promise { + const endpoints = await discoverDatabricksOidcEndpoints(input.workspaceHost); + return { + grantType: "authorization-code", + credentialBaseUrl: normalizeDatabricksWorkspaceHost(input.workspaceHost), + clientId: input.clientId, + scope: input.scope, + authorizationEndpoint: endpoints.authorizationEndpoint, + tokenEndpoint: endpoints.tokenEndpoint, + receiver: input.receiver, + ...(input.timeoutMs ? { timeoutMs: input.timeoutMs } : {}), + }; +} + +export async function createDatabricksClientCredentialsGrant(input: { + workspaceHost: string; + clientId: string; + clientSecret: string; + scope?: string; +}): Promise { + const host = normalizeDatabricksWorkspaceHost(input.workspaceHost); + const endpoints = await discoverDatabricksOidcEndpoints(host); + const secretFingerprint = await sha256(input.clientSecret); + return { + grantType: "client-credentials", + credentialBaseUrl: host, + clientId: input.clientId, + clientSecret: input.clientSecret, + scope: input.scope ?? "all-apis", + tokenEndpoint: endpoints.tokenEndpoint, + cacheKey: `${host}\n${input.clientId}\n${secretFingerprint}`, + }; +} + +function hasStringProperty( + value: unknown, + key: Key, +): value is Record { + return ( + typeof value === "object" && + value !== null && + key in value && + typeof Reflect.get(value, key) === "string" + ); +} + +async function sha256(value: string): Promise { + const digest = await crypto.subtle.digest("SHA-256", new TextEncoder().encode(value)); + return Array.from(new Uint8Array(digest), (byte) => byte.toString(16).padStart(2, "0")).join(""); +} + +function validateEndpoint(workspaceHost: string, raw: string): string { + const endpoint = new URL(raw, workspaceHost); + const workspace = new URL(workspaceHost); + if (endpoint.origin !== workspace.origin) { + throw new Error("Databricks OIDC discovery returned a cross-origin endpoint"); + } + return endpoint.toString(); +} diff --git a/packages/ai-credentials/src/index.ts b/packages/ai-credentials/src/index.ts index ce84734..03e5573 100644 --- a/packages/ai-credentials/src/index.ts +++ b/packages/ai-credentials/src/index.ts @@ -15,10 +15,40 @@ * (vscode) and are injected here. */ -export type { CredentialProvider, Disposable } from "./CredentialProvider"; -export type { Backend, OAuthBackendHooks, OAuthProviderConfig, StoredOAuthTokens } from "./Backend"; +export type { + AuthenticationChallenge, + AuthenticationStartResult, + CredentialMutation, + CredentialProvider, + CredentialSourceInput, + CredentialStatus, + Disposable, + MutableCredentialProvider, +} from "./CredentialProvider"; +export type { + AcquisitionBackendHooks, + AuthenticationCommitResult, + AuthorizationCodeCallback, + AuthorizationCodeReceiver, + Backend, + CredentialSourceContext, + MutableBackend, + OAuthBackendHooks, + OAuthGrantConfig, + OAuthProviderConfig, + PreparedAuthorizationCodeReceiver, + StoredOAuthTokens, +} from "./Backend"; export { createCredentialProvider } from "./createCredentialProvider"; export type { CreateCredentialProviderOptions, CredentialProviderHandle, + MutableCredentialProviderHandle, } from "./createCredentialProvider"; +export { + createDatabricksAuthorizationCodeGrant, + createDatabricksClientCredentialsGrant, + discoverDatabricksOidcEndpoints, + normalizeDatabricksWorkspaceHost, +} from "./databricks-oauth"; +export type { DatabricksOidcEndpoints } from "./databricks-oauth"; diff --git a/packages/ai-credentials/src/positron/PositronBackend.ts b/packages/ai-credentials/src/positron/PositronBackend.ts index d490e24..fc503d9 100644 --- a/packages/ai-credentials/src/positron/PositronBackend.ts +++ b/packages/ai-credentials/src/positron/PositronBackend.ts @@ -77,6 +77,14 @@ export function createVscodeCredentialConfig(): CredentialConfig { account: snowflakeConfig?.SNOWFLAKE_ACCOUNT || process.env.SNOWFLAKE_ACCOUNT, }; }, + getDatabricks: () => { + const databricksConfig = vscode.workspace + .getConfiguration("authentication.databricks") + .get>("credentials"); + return { + host: databricksConfig?.DATABRICKS_HOST || process.env.DATABRICKS_HOST, + }; + }, }; } diff --git a/packages/ai-credentials/src/positron/__tests__/PositronBackend.test.ts b/packages/ai-credentials/src/positron/__tests__/PositronBackend.test.ts index cd6eb83..d583509 100644 --- a/packages/ai-credentials/src/positron/__tests__/PositronBackend.test.ts +++ b/packages/ai-credentials/src/positron/__tests__/PositronBackend.test.ts @@ -63,6 +63,7 @@ const PROVIDER_MAP: Record = { fallbackScopes: [["read:user", "user:email", "repo", "workflow"], ["user:email"]], credentialType: "apikey", }, + databricks: { authProviderId: "databricks", scopes: [], credentialType: "apikey" }, }; const logger = { info: vi.fn(), warn: vi.fn(), error: vi.fn(), debug: vi.fn(), trace: vi.fn() }; @@ -140,6 +141,100 @@ describe("createPositronBackend", () => { expect(configChangeHook.callback).toBeNull(); }); + // --- Databricks host resolution (ported from the removed bridge auth suite) --- + + it("resolves the Databricks base URL from authentication.databricks.credentials.DATABRICKS_HOST", async () => { + mockGetSession.mockResolvedValue(makeSession("databricks-bearer-token")); + mockGetConfigValue.mockImplementation((key: string) => { + if (key === "credentials") { + return { DATABRICKS_HOST: "https://adb-123.4.azuredatabricks.net" }; + } + return undefined; + }); + const backend = createPositronBackend({ logger, providerMap: PROVIDER_MAP }); + + expect(await backend.getCredentials("databricks")).toEqual({ + type: "apikey", + apiKey: "databricks-bearer-token", + baseUrl: "https://adb-123.4.azuredatabricks.net", + customHeaders: undefined, + }); + expect(mockGetSession).toHaveBeenCalledWith("databricks", [], { silent: true }); + }); + + it("normalizes a scheme-less Databricks host with trailing slash", async () => { + mockGetSession.mockResolvedValue(makeSession("databricks-bearer-token")); + mockGetConfigValue.mockImplementation((key: string) => { + if (key === "credentials") { + return { DATABRICKS_HOST: "my-workspace.cloud.databricks.com/" }; + } + return undefined; + }); + const backend = createPositronBackend({ logger, providerMap: PROVIDER_MAP }); + + expect(await backend.getCredentials("databricks")).toMatchObject({ + type: "apikey", + baseUrl: "https://my-workspace.cloud.databricks.com", + }); + }); + + it("falls back to the DATABRICKS_HOST env var when the setting is absent", async () => { + mockGetSession.mockResolvedValue(makeSession("databricks-bearer-token")); + const backend = createPositronBackend({ logger, providerMap: PROVIDER_MAP }); + + const originalHost = process.env.DATABRICKS_HOST; + process.env.DATABRICKS_HOST = "https://env-workspace.cloud.databricks.com"; + try { + expect(await backend.getCredentials("databricks")).toMatchObject({ + baseUrl: "https://env-workspace.cloud.databricks.com", + }); + } finally { + if (originalHost === undefined) { + delete process.env.DATABRICKS_HOST; + } else { + process.env.DATABRICKS_HOST = originalHost; + } + } + }); + + it("leaves the Databricks base URL undefined when no host is configured", async () => { + mockGetSession.mockResolvedValue(makeSession("databricks-bearer-token")); + const backend = createPositronBackend({ logger, providerMap: PROVIDER_MAP }); + + const originalHost = process.env.DATABRICKS_HOST; + delete process.env.DATABRICKS_HOST; + try { + expect(await backend.getCredentials("databricks")).toMatchObject({ + baseUrl: undefined, + }); + } finally { + if (originalHost !== undefined) { + process.env.DATABRICKS_HOST = originalHost; + } + } + }); + + it("reads Databricks customHeaders alongside the host", async () => { + mockGetSession.mockResolvedValue(makeSession("databricks-bearer-token")); + mockGetConfigValue.mockImplementation((key: string) => { + if (key === "credentials") { + return { DATABRICKS_HOST: "https://adb-123.4.azuredatabricks.net" }; + } + if (key === "databricks.customHeaders") { + return { "x-databricks-use-coding-agent-mode": "true" }; + } + return undefined; + }); + const backend = createPositronBackend({ logger, providerMap: PROVIDER_MAP }); + + expect(await backend.getCredentials("databricks")).toEqual({ + type: "apikey", + apiKey: "databricks-bearer-token", + baseUrl: "https://adb-123.4.azuredatabricks.net", + customHeaders: { "x-databricks-use-coding-agent-mode": "true" }, + }); + }); + // --- Copilot fallback scopes (ported from the removed bridge auth suite) --- it("uses the primary read:user scope for copilot when a session exists there", async () => { diff --git a/packages/ai-credentials/src/store-backend/StoreBackend.ts b/packages/ai-credentials/src/store-backend/StoreBackend.ts index 66e917c..2565f3b 100644 --- a/packages/ai-credentials/src/store-backend/StoreBackend.ts +++ b/packages/ai-credentials/src/store-backend/StoreBackend.ts @@ -2,232 +2,716 @@ * Copyright (C) 2026 Posit Software, PBC. All rights reserved. *--------------------------------------------------------------------------------------------*/ -/** - * Store-backed credential Backend. - * - * Reads runtime credentials from the generic {@link SingleFileStore} with an - * environment-variable fallback (`store → env → null`), and supplies the OAuth - * device-flow hooks that the root state machine calls (option B). It owns the - * on-disk {@link StoredProviderCredentials} shape (guarded by a tolerant Zod - * schema) and the storage-key scheme via {@link storageKeyFor}. - * - * PURITY: imports `ai-credentials/store` (fs) + `ai-credentials/types` (pure) - * and the root {@link Backend} type — but nothing from `@assistant/*`, so - * standalone consumers (Notebooks) resolve credentials without the assistant - * monorepo. - * - * Provider descriptors (which auth method a provider uses) are **injected** via - * `resolveAuthMethod`, and OAuth connection config via `oauthConfigForProvider`, - * so this backend needs neither the provider registry nor the catalog. - */ - import type { - Backend, + AcquisitionBackendHooks, + CredentialSourceContext, + MutableBackend, OAuthBackendHooks, + OAuthGrantConfig, OAuthProviderConfig, StoredOAuthTokens, } from "../Backend"; -import type { Disposable } from "../CredentialProvider"; +import type { + CredentialMutation, + CredentialSourceInput, + CredentialStatus, + Disposable, +} from "../CredentialProvider"; import type { SingleFileStore } from "../store"; import type { Logger, ProviderCredentials, TokenData } from "../types"; -import { storageKeyFor } from "../types"; +import { normalizeDatabricksHost, storageKeyFor } from "../types"; import { resolveCredentialsFromEnv } from "./envCredentialResolver"; import { storedProviderCredentialsSchema, type StoredProviderCredentials, } from "./StoredProviderCredentials"; -/** Auth descriptor for a provider (injected — derived from registry/catalog). */ export interface AuthMethodDescriptor { authMethodId: string; apiKeyOptional?: boolean; } export interface CreateStoreBackendOptions { - /** The generic single-file credential store. */ store: SingleFileStore; - /** - * Resolve a provider's auth method + flags. Returns undefined for unknown - * providers. Injected so the backend needs no registry/catalog import. - */ resolveAuthMethod(providerId: string): AuthMethodDescriptor | undefined; - /** - * OAuth connection config for device-flow providers (e.g. positai), or - * undefined for providers/hosts without device-flow support. When omitted - * entirely, the backend exposes no OAuth hooks. - */ - oauthConfigForProvider?: (providerId: string) => OAuthProviderConfig | undefined; - /** Called when a provider's credentials become ready/refreshed. */ + oauthConfigForProvider?: ( + providerId: string, + source?: CredentialSourceContext, + ) => + | OAuthGrantConfig + | OAuthProviderConfig + | undefined + | Promise; + shapeToken?: ( + providerId: string, + accessToken: string, + config: OAuthGrantConfig, + source: CredentialSourceContext, + ) => ProviderCredentials; notifyReady?: (providerId: string) => void; - /** - * Provider ids whose store records should trigger `onDidChangeCredentials` - * on file change. Empty/omitted → the change subscription is a no-op. - */ watchedProviderIds?: string[]; - /** Environment variables for the env fallback (defaults to `process.env`). */ env?: Record; logger?: Logger; + generationFactory?: () => string; +} + +interface NormalizedStored { + readiness: "pending" | "ready" | "unauthenticated"; + generation?: string; + source?: CredentialSourceInput; + tokens?: StoredOAuthTokens; + error?: string; } const NOOP_DISPOSABLE: Disposable = { dispose() {} }; -/** Build a store-backed {@link Backend}. */ -export function createStoreBackend(options: CreateStoreBackendOptions): Backend { +export function createStoreBackend(options: CreateStoreBackendOptions): MutableBackend { const { store, resolveAuthMethod, oauthConfigForProvider, + shapeToken, notifyReady, - watchedProviderIds, - env, + watchedProviderIds = [], logger, + generationFactory = defaultGeneration, } = options; + const currentEnvironment = () => options.env ?? process.env; + + function keyFor(providerId: string): string | undefined { + const descriptor = resolveAuthMethod(providerId); + return descriptor ? storageKeyFor(providerId, descriptor.authMethodId) : undefined; + } - /** - * Read a credential record and validate it against the tolerant Zod schema - * (Phase 0 #5 runtime guard). Legacy records — which populate only a subset - * of fields — parse unchanged; structurally invalid records (e.g. an - * `apiKeyAuth` missing its required `apiKey`) are dropped rather than flowing - * out as credentials or into OAuth refresh. Returns undefined when the key is - * absent or the record fails validation. - */ - async function readRecord(key: string): Promise { + async function readRecord(providerId: string): Promise { + const key = keyFor(providerId); + if (!key) return undefined; const raw = await store.get(key); if (raw === undefined) return undefined; const parsed = storedProviderCredentialsSchema.safeParse(raw); if (!parsed.success) { - logger?.warn( - `[ai-credentials/store-backend] Ignoring malformed credential record at ${key}: ${parsed.error.message}`, - ); + logger?.warn(`[ai-credentials] Ignoring invalid credential record for ${providerId}`); return undefined; } return parsed.data; } - async function getCredentials(providerId: string): Promise { - const descriptor = resolveAuthMethod(providerId); - if (!descriptor) return null; - const { authMethodId, apiKeyOptional } = descriptor; - - // OAuth is resolved by the root state machine via the oauth hooks, not here. - if (authMethodId === "oauth") return null; - - const stored = await readRecord(storageKeyFor(providerId, authMethodId)); - - if (stored) { - if (authMethodId === "apikey" && stored.apiKeyAuth) { - // A stored-but-empty required key falls through to the env fallback. - if (apiKeyOptional || stored.apiKeyAuth.apiKey) { - return { - type: "apikey", - apiKey: stored.apiKeyAuth.apiKey, - baseUrl: stored.apiKeyAuth.baseUrl, - }; - } - } + function normalize(record: StoredProviderCredentials | undefined): NormalizedStored | null { + if (!record) return null; + const readiness = + record.readiness ?? + (record.authenticated === true + ? "ready" + : record.authenticated === false + ? "unauthenticated" + : undefined); - if (authMethodId === "local" && stored.localAuth?.endpoint) { - return { type: "local", endpoint: stored.localAuth.endpoint }; - } + if (record.source === "oauth-m2m" && record.clientCredentialsAuth) { + return { + readiness: readiness === "pending" ? "pending" : "ready", + generation: record.generation, + source: { type: "oauth-m2m", ...record.clientCredentialsAuth }, + }; + } + + if (record.source === "oauth-u2m") { + const workspaceHost = record.oauthAuth?.workspaceHost; + if (!workspaceHost) + return { readiness: "unauthenticated", generation: record.generation, error: record.error }; + return { + readiness: readiness ?? "unauthenticated", + generation: record.generation, + source: { type: "oauth-u2m", workspaceHost }, + tokens: readiness === "ready" ? storedTokens(record) : undefined, + error: record.error, + }; + } + + if (record.apiKeyAuth && (record.source === undefined || record.source === "api-key")) { + return { + readiness: "ready", + generation: record.generation, + source: { + type: "api-key", + apiKey: record.apiKeyAuth.apiKey, + ...(record.apiKeyAuth.baseUrl ? { baseUrl: record.apiKeyAuth.baseUrl } : {}), + }, + }; + } + + if (record.source === "oauth-device" || (record.source === undefined && record.oauthAuth)) { + const tokenData = record.oauthAuth?.tokenData; + return { + readiness: readiness ?? (tokenData ? "ready" : "unauthenticated"), + generation: record.generation, + source: { type: "oauth-device" }, + tokens: + readiness === "ready" || (readiness === undefined && tokenData) + ? storedTokens(record) + : undefined, + error: record.error, + }; + } + + if (record.localAuth) { + return { + readiness: "ready", + generation: record.generation, + source: { type: "local", endpoint: record.localAuth.endpoint }, + }; + } + if (record.awsAuth) { + return { + readiness: "ready", + generation: record.generation, + source: { type: "aws-credentials", ...record.awsAuth }, + }; + } + if (record.googleCloudAuth) { + return { + readiness: "ready", + generation: record.generation, + source: { type: "google-cloud", ...record.googleCloudAuth }, + }; + } + return { readiness: "unauthenticated", generation: record.generation, error: record.error }; + } + + async function storedSource(providerId: string): Promise { + return normalize(await readRecord(providerId)); + } + + type EnvironmentResolution = + | { + kind: "credentials"; + credentials: ProviderCredentials; + } + | { + kind: "oauth-m2m"; + source: Extract; + } + | { + kind: "incomplete-oauth-m2m"; + workspaceHost?: string; + error: string; + } + | { kind: "none" }; + + function environmentResolution(providerId: string): EnvironmentResolution { + const env = currentEnvironment(); + if (providerId !== "databricks") { + const credentials = resolveCredentialsFromEnv(providerId, env); + return credentials ? { kind: "credentials", credentials } : { kind: "none" }; + } - if (authMethodId === "aws-credentials" && stored.awsAuth?.region) { + const token = env.DATABRICKS_TOKEN; + const explicitlyM2m = env.DATABRICKS_AUTH_TYPE === "oauth-m2m"; + if (token && !explicitlyM2m) { + const credentials = resolveCredentialsFromEnv(providerId, env); + return credentials ? { kind: "credentials", credentials } : { kind: "none" }; + } + if (env.DATABRICKS_CLIENT_ID && env.DATABRICKS_CLIENT_SECRET && env.DATABRICKS_HOST) { + let workspaceHost: string; + try { + workspaceHost = normalizeDatabricksHost(env.DATABRICKS_HOST); + } catch (error) { return { - type: "aws-credentials", - region: stored.awsAuth.region, - profile: stored.awsAuth.profile, - accessKeyId: stored.awsAuth.accessKeyId, - secretAccessKey: stored.awsAuth.secretAccessKey, - sessionToken: stored.awsAuth.sessionToken, + kind: "incomplete-oauth-m2m", + workspaceHost: env.DATABRICKS_HOST, + error: error instanceof Error ? error.message : "Invalid Databricks workspace URL", }; } + return { + kind: "oauth-m2m", + source: { + type: "oauth-m2m", + origin: "environment", + clientId: env.DATABRICKS_CLIENT_ID, + clientSecret: env.DATABRICKS_CLIENT_SECRET, + workspaceHost, + }, + }; + } + if (explicitlyM2m) { + return { + kind: "incomplete-oauth-m2m", + workspaceHost: env.DATABRICKS_HOST, + error: + "Databricks OAuth M2M requires DATABRICKS_HOST, DATABRICKS_CLIENT_ID, and DATABRICKS_CLIENT_SECRET", + }; + } + return { kind: "none" }; + } - if (authMethodId === "google-cloud" && stored.googleCloudAuth?.project) { - return { - type: "google-cloud", - project: stored.googleCloudAuth.project, - location: stored.googleCloudAuth.location, - }; + async function sourceContext(providerId: string): Promise { + const normalized = await storedSource(providerId); + if (normalized?.source) { + if (normalized.source.type === "oauth-device") { + return { type: "oauth-device", origin: "stored" }; + } + if (normalized.source.type === "oauth-u2m") { + return { ...normalized.source, origin: "stored" }; + } + if (normalized.source.type === "oauth-m2m") { + return { ...normalized.source, origin: "stored" }; + } + return undefined; + } + if (providerId === "databricks") { + const environment = environmentResolution(providerId); + return environment.kind === "oauth-m2m" ? environment.source : undefined; + } + if (resolveAuthMethod(providerId)?.authMethodId === "oauth") { + return { type: "oauth-device", origin: "implicit" }; + } + return undefined; + } + + async function resolveGrant(providerId: string): Promise { + if (!oauthConfigForProvider) return undefined; + const source = await sourceContext(providerId); + if (!source) return undefined; + const config = await oauthConfigForProvider(providerId, source); + if (!config) return undefined; + if ("grantType" in config) return config; + return { + grantType: "device-code", + clientId: config.clientId, + scope: config.scope, + deviceAuthorizationEndpoint: `https://${config.authHost}/oauth/device/authorize`, + tokenEndpoint: `https://${config.authHost}/oauth/token`, + }; + } + + async function getCredentials(providerId: string): Promise { + const descriptor = resolveAuthMethod(providerId); + if (!descriptor) return null; + const normalized = await storedSource(providerId); + if (normalized?.source) { + const source = normalized.source; + switch (source.type) { + case "api-key": + if (!source.apiKey && !descriptor.apiKeyOptional) break; + return { type: "apikey", apiKey: source.apiKey, baseUrl: source.baseUrl }; + case "local": + return { type: "local", endpoint: source.endpoint }; + case "aws-credentials": + return { ...source }; + case "google-cloud": + return { ...source }; + case "oauth-device": + case "oauth-u2m": + case "oauth-m2m": + return null; } } + const environment = environmentResolution(providerId); + return environment.kind === "credentials" ? environment.credentials : null; + } + + async function mutateCredentials( + providerId: string, + mutation: CredentialMutation, + ): Promise { + const key = keyFor(providerId); + if (!key) throw new Error(`Unknown provider: ${providerId}`); + await store.withLock(async () => { + await store.get(key); + const generation = generationFactory(); + if (mutation.kind === "clear") { + await store.set(key, { + generation, + readiness: "unauthenticated", + configured: false, + authenticated: false, + } satisfies StoredProviderCredentials); + return; + } + await store.set(key, recordForSource(mutation.source, generation)); + }); + } - // Fall back to environment variables (secret fields only). - const envCredentials = resolveCredentialsFromEnv(providerId, env); - if (envCredentials) { - logger?.debug(`[ai-credentials/store-backend] Using env credentials for ${providerId}`); - return envCredentials; + async function getCredentialSource(providerId: string): Promise { + const normalized = await storedSource(providerId); + return normalized?.source ?? null; + } + + async function getCredentialStatus(providerId: string): Promise { + const normalized = await storedSource(providerId); + if (normalized?.source) { + const metadata = sourceMetadata(normalized.source); + return { + configured: true, + authenticated: normalized.readiness === "ready", + readiness: normalized.readiness, + source: normalized.source.type, + origin: "stored", + expiresAt: normalized.tokens?.expiresAt, + scope: normalized.tokens?.scope, + error: "error" in normalized ? normalized.error : undefined, + metadata, + }; + } + const environment = environmentResolution(providerId); + if (environment.kind === "oauth-m2m") { + return { + configured: true, + authenticated: true, + readiness: "ready", + source: "oauth-m2m", + origin: "environment", + metadata: { workspaceHost: environment.source.workspaceHost }, + }; } + if (environment.kind === "credentials") { + return { + configured: true, + authenticated: true, + readiness: "ready", + source: "api-key", + origin: "environment", + }; + } + if (environment.kind === "incomplete-oauth-m2m") { + return { + configured: false, + authenticated: false, + readiness: "unauthenticated", + source: "oauth-m2m", + origin: "environment", + error: environment.error, + metadata: environment.workspaceHost + ? { workspaceHost: environment.workspaceHost } + : undefined, + }; + } + return { + configured: false, + authenticated: false, + readiness: "unauthenticated", + error: normalized?.error, + }; + } - return null; + async function beginAuthentication(providerId: string): Promise { + const key = keyFor(providerId); + if (!key) throw new Error(`Unknown provider: ${providerId}`); + return store.withLock(async () => { + const normalized = normalize(await readRecord(providerId)); + const generation = generationFactory(); + let source: CredentialSourceInput = { type: "oauth-device" }; + if (normalized?.source) source = normalized.source; + if (source.type !== "oauth-device" && source.type !== "oauth-u2m") { + throw new Error(`Stored source ${source.type} is not interactive`); + } + await store.set(key, pendingRecord(source, generation)); + return generation; + }); } - function onDidChangeCredentials(callback: (providerIds: string[]) => void): Disposable { - if (!watchedProviderIds || watchedProviderIds.length === 0) { - return NOOP_DISPOSABLE; + async function commitAuthentication( + providerId: string, + generation: string, + tokens: TokenData, + ): Promise<"committed" | "superseded"> { + return compareAndWrite(providerId, generation, (current) => { + if (!current.source) return null; + if (current.source.type !== "oauth-device" && current.source.type !== "oauth-u2m") + return null; + return authenticatedOAuthRecord(current.source, tokens, generationFactory()); + }); + } + + async function finishAuthentication( + providerId: string, + generation: string, + error: string, + ): Promise<"committed" | "superseded"> { + return compareAndWrite(providerId, generation, (current) => { + if (!current.source) return null; + return terminalOAuthRecord(current.source, generationFactory(), error); + }); + } + + async function compareAndWrite( + providerId: string, + generation: string, + build: (current: NormalizedStored) => StoredProviderCredentials | null, + ): Promise<"committed" | "superseded"> { + const key = keyFor(providerId); + if (!key) return "superseded"; + return store.withLock(async () => { + const current = normalize(await readRecord(providerId)); + if (!current || current.generation !== generation) return "superseded"; + const next = build(current); + if (!next) return "superseded"; + await store.set(key, next); + return "committed"; + }); + } + + async function readTokens(providerId: string): Promise { + const normalized = await storedSource(providerId); + if (!normalized || normalized.readiness !== "ready" || !normalized.source) return null; + if (normalized.source.type !== "oauth-device" && normalized.source.type !== "oauth-u2m") + return null; + return normalized.tokens ?? null; + } + + async function persistRefreshedTokens(providerId: string, tokens: TokenData): Promise { + const key = keyFor(providerId); + if (!key) return; + const current = normalize(await readRecord(providerId)); + if (!current?.source) return; + if (current.source.type !== "oauth-device" && current.source.type !== "oauth-u2m") return; + await store.set(key, authenticatedOAuthRecord(current.source, tokens, generationFactory())); + } + + async function persistRefreshError(providerId: string, error: string): Promise { + const key = keyFor(providerId); + if (!key) return; + const current = normalize(await readRecord(providerId)); + if (!current?.source) return; + await store.set(key, terminalOAuthRecord(current.source, generationFactory(), error)); + } + + const acquisition: AcquisitionBackendHooks | undefined = oauthConfigForProvider + ? { + configForProvider: resolveGrant, + readTokens, + beginAuthentication, + commitAuthentication, + finishAuthentication, + persistRefreshedTokens, + persistRefreshError, + withRefreshTransaction: (_providerId, operation) => store.withLock(operation), + shapeToken: asyncShapeToken, + notifyReady(providerId) { + notifyReady?.(providerId); + }, + } + : undefined; + + function asyncShapeToken( + providerId: string, + accessToken: string, + config: OAuthGrantConfig, + ): ProviderCredentials { + const descriptor = resolveAuthMethod(providerId); + // shapeToken is intentionally synchronous. The source was already resolved + // by configForProvider, so provider-specific shaping can be derived from the + // descriptor/config. Databricks baseUrl is attached by Node's catalog merge. + if (shapeToken) { + // Custom shapers that need the source should encode its non-secret identity + // in the resolved grant configuration. Keep secrets out of errors/logs. + const fallback: CredentialSourceContext = { type: "oauth-device", origin: "implicit" }; + return shapeToken(providerId, accessToken, config, fallback); } - const ids = [...watchedProviderIds]; - return store.watch(() => callback(ids)); + return descriptor?.authMethodId === "apikey" + ? { type: "apikey", apiKey: accessToken, baseUrl: config.credentialBaseUrl } + : { type: "oauth", accessToken }; } - // --- OAuth hooks (only when device-flow config can be supplied) ----------- - let oauth: OAuthBackendHooks | undefined; - if (oauthConfigForProvider) { - const oauthKey = (providerId: string): string => storageKeyFor(providerId, "oauth"); + const oauth: OAuthBackendHooks | undefined = oauthConfigForProvider + ? { + configForProvider(providerId): OAuthProviderConfig | undefined { + // Compatibility is limited to the original synchronous Posit AI callback. + const value = oauthConfigForProvider(providerId, { + type: "oauth-device", + origin: "implicit", + }); + if (value instanceof Promise || !value || "grantType" in value) return undefined; + return value; + }, + readTokens, + async persistTokens(providerId, tokens) { + const key = keyFor(providerId); + if (!key) return; + await store.withLock(async () => { + await store.set( + key, + authenticatedOAuthRecord({ type: "oauth-device" }, tokens, generationFactory()), + ); + }); + }, + persistError: persistRefreshError, + async clearError(providerId) { + const key = keyFor(providerId); + if (!key) return; + await store.withLock(async () => { + await store.set( + key, + terminalOAuthRecord({ type: "oauth-device" }, generationFactory()), + ); + }); + }, + notifyReady(providerId) { + notifyReady?.(providerId); + }, + } + : undefined; - oauth = { - configForProvider: oauthConfigForProvider, + function onDidChangeCredentials(callback: (providerIds: string[]) => void): Disposable { + if (watchedProviderIds.length === 0) return NOOP_DISPOSABLE; + return store.watch(() => callback([...watchedProviderIds])); + } - async readTokens(providerId: string): Promise { - const stored = await readRecord(oauthKey(providerId)); - // Gate on authenticated: a refresh-failure record can keep stale - // `oauthAuth.tokenData` while marked `authenticated: false`; those - // tokens must not be treated as usable material. - if (!stored || stored.authenticated !== true) return null; - const tokenData = stored.oauthAuth?.tokenData; - if (!tokenData) return null; - return { - accessToken: tokenData.accessToken, - refreshToken: tokenData.refreshToken, - expiresAt: tokenData.expiresAt, - scope: tokenData.scope, - tokenType: tokenData.tokenType, - }; - }, + return { + getCredentials, + onDidChangeCredentials, + ...(oauth ? { oauth } : {}), + ...(acquisition ? { acquisition } : {}), + mutateCredentials, + getCredentialStatus, + getCredentialSource, + }; +} - async persistTokens(providerId: string, tokens: TokenData): Promise { - const expiresAt = new Date(Date.now() + tokens.expiresIn * 1000).toISOString(); - const record: StoredProviderCredentials = { - authenticated: true, - oauthAuth: { - tokenData: { - accessToken: tokens.accessToken, - refreshToken: tokens.refreshToken, - expiresAt, - tokenType: tokens.tokenType, - scope: tokens.scope, - }, - expiresAt, - scope: tokens.scope, - }, - error: undefined, - }; - await store.set(oauthKey(providerId), record); - }, +function storedTokens(record: StoredProviderCredentials): StoredOAuthTokens | undefined { + const tokenData = record.oauthAuth?.tokenData; + if (!tokenData) return undefined; + return { + accessToken: tokenData.accessToken, + refreshToken: tokenData.refreshToken, + expiresAt: tokenData.expiresAt, + tokenType: tokenData.tokenType, + scope: tokenData.scope, + }; +} - async persistError(providerId: string, error: string): Promise { - await store.set(oauthKey(providerId), { authenticated: false, error }); - }, +function recordForSource( + source: CredentialSourceInput, + generation: string, +): StoredProviderCredentials { + switch (source.type) { + case "api-key": + return { + generation, + readiness: "ready", + source: "api-key", + configured: true, + authenticated: true, + apiKeyAuth: { + apiKey: source.apiKey, + ...(source.baseUrl ? { baseUrl: source.baseUrl } : {}), + }, + }; + case "oauth-device": + return terminalOAuthRecord(source, generation); + case "oauth-u2m": + return terminalOAuthRecord(source, generation); + case "oauth-m2m": + return { + generation, + readiness: "ready", + source: "oauth-m2m", + configured: true, + authenticated: true, + clientCredentialsAuth: { + clientId: source.clientId, + clientSecret: source.clientSecret, + workspaceHost: source.workspaceHost, + }, + }; + case "local": + return { + generation, + readiness: "ready", + configured: true, + localAuth: { endpoint: source.endpoint }, + }; + case "aws-credentials": { + const { type: _type, ...awsAuth } = source; + return { generation, readiness: "ready", configured: true, awsAuth }; + } + case "google-cloud": + return { + generation, + readiness: "ready", + configured: true, + googleCloudAuth: { project: source.project, location: source.location }, + }; + } +} - async clearError(providerId: string): Promise { - await store.set(oauthKey(providerId), { authenticated: false, error: undefined }); - }, +function pendingRecord( + source: Extract, + generation: string, +): StoredProviderCredentials { + return { + generation, + readiness: "pending", + source: source.type, + configured: true, + authenticated: false, + oauthAuth: source.type === "oauth-u2m" ? { workspaceHost: source.workspaceHost } : undefined, + }; +} - notifyReady(providerId: string): void { - notifyReady?.(providerId); +function authenticatedOAuthRecord( + source: Extract, + tokens: TokenData, + generation: string, +): StoredProviderCredentials { + const expiresAt = new Date(Date.now() + tokens.expiresIn * 1000).toISOString(); + return { + generation, + readiness: "ready", + source: source.type, + configured: true, + authenticated: true, + oauthAuth: { + tokenData: { + accessToken: tokens.accessToken, + refreshToken: tokens.refreshToken, + expiresAt, + tokenType: tokens.tokenType, + scope: tokens.scope, }, + expiresAt, + scope: tokens.scope, + ...(source.type === "oauth-u2m" ? { workspaceHost: source.workspaceHost } : {}), + }, + }; +} + +function terminalOAuthRecord( + source: CredentialSourceInput, + generation: string, + error?: string, +): StoredProviderCredentials { + if (source.type !== "oauth-device" && source.type !== "oauth-u2m") { + return { + generation, + readiness: "unauthenticated", + configured: false, + authenticated: false, + error, }; } + return { + generation, + readiness: "unauthenticated", + source: source.type, + configured: true, + authenticated: false, + error, + oauthAuth: source.type === "oauth-u2m" ? { workspaceHost: source.workspaceHost } : undefined, + }; +} + +function sourceMetadata(source: CredentialSourceInput): Record | undefined { + if (source.type === "api-key" && source.baseUrl) return { baseUrl: source.baseUrl }; + if (source.type === "oauth-u2m" || source.type === "oauth-m2m") { + return { workspaceHost: source.workspaceHost }; + } + if (source.type === "local") return { endpoint: source.endpoint }; + if (source.type === "google-cloud") return { project: source.project, location: source.location }; + return undefined; +} - return oauth - ? { getCredentials, onDidChangeCredentials, oauth } - : { getCredentials, onDidChangeCredentials }; +function defaultGeneration(): string { + const bytes = new Uint8Array(16); + globalThis.crypto.getRandomValues(bytes); + return [...bytes].map((value) => value.toString(16).padStart(2, "0")).join(""); } diff --git a/packages/ai-credentials/src/store-backend/StoredProviderCredentials.ts b/packages/ai-credentials/src/store-backend/StoredProviderCredentials.ts index 4b5ff2f..db51c17 100644 --- a/packages/ai-credentials/src/store-backend/StoredProviderCredentials.ts +++ b/packages/ai-credentials/src/store-backend/StoredProviderCredentials.ts @@ -34,6 +34,8 @@ const tokenDataSchema = z.object({ scope: z.string(), }); +const credentialSourceSchema = z.enum(["api-key", "oauth-device", "oauth-u2m", "oauth-m2m"]); + /** * Tolerant Zod schema for a single credential record in data.json. * @@ -48,6 +50,9 @@ export const storedProviderCredentialsSchema = z.object({ authenticated: z.boolean().optional(), error: z.string().optional(), metadata: z.record(z.string(), z.unknown()).optional(), + generation: z.string().optional(), + readiness: z.enum(["pending", "ready", "unauthenticated"]).optional(), + source: credentialSourceSchema.optional(), // Auth-method-specific fields (only one should be populated) apiKeyAuth: z @@ -59,9 +64,18 @@ export const storedProviderCredentialsSchema = z.object({ oauthAuth: z .object({ - tokenData: tokenDataSchema, - expiresAt: z.string(), - scope: z.string(), + tokenData: tokenDataSchema.optional(), + expiresAt: z.string().optional(), + scope: z.string().optional(), + workspaceHost: z.string().optional(), + }) + .optional(), + + clientCredentialsAuth: z + .object({ + clientId: z.string(), + clientSecret: z.string(), + workspaceHost: z.string(), }) .optional(), diff --git a/packages/ai-credentials/src/store-backend/__tests__/StoreBackend.test.ts b/packages/ai-credentials/src/store-backend/__tests__/StoreBackend.test.ts index a31fc01..a0c0d93 100644 --- a/packages/ai-credentials/src/store-backend/__tests__/StoreBackend.test.ts +++ b/packages/ai-credentials/src/store-backend/__tests__/StoreBackend.test.ts @@ -27,6 +27,7 @@ const DESCRIPTORS: Record = { ollama: { authMethodId: "local" }, bedrock: { authMethodId: "aws-credentials" }, positai: { authMethodId: "oauth" }, + databricks: { authMethodId: "apikey" }, }; function resolveAuthMethod(id: string): AuthMethodDescriptor | undefined { @@ -241,6 +242,10 @@ describe("createStoreBackend", () => { expect(await oauth.readTokens("positai")).toBeNull(); const stored = await store.get(storageKeyFor("positai", "oauth")); expect(stored?.error).toBe("access_denied"); + expect(await backend.getCredentialStatus("positai")).toMatchObject({ + authenticated: false, + error: "access_denied", + }); }); it("clearError resets a prior error record to a clean unauthenticated state", async () => { @@ -290,6 +295,84 @@ describe("createStoreBackend", () => { }); }); + describe("Databricks environment source selection", () => { + it("rejects an insecure environment M2M workspace without falling back", async () => { + const backend = createStoreBackend({ + store, + resolveAuthMethod, + oauthConfigForProvider: () => undefined, + env: { + DATABRICKS_AUTH_TYPE: "oauth-m2m", + DATABRICKS_HOST: "http://workspace.test", + DATABRICKS_CLIENT_ID: "client", + DATABRICKS_CLIENT_SECRET: "secret", + }, + }); + + expect(await backend.getCredentials("databricks")).toBeNull(); + expect(await backend.getCredentialStatus("databricks")).toMatchObject({ + configured: false, + authenticated: false, + error: "Databricks workspace URL must use HTTPS", + }); + }); + + it("does not report PAT readiness when explicitly selected M2M is incomplete", async () => { + const backend = createStoreBackend({ + store, + resolveAuthMethod, + env: { + DATABRICKS_AUTH_TYPE: "oauth-m2m", + DATABRICKS_TOKEN: "must-not-be-used", + DATABRICKS_HOST: "https://workspace.test", + DATABRICKS_CLIENT_ID: "incomplete-client", + }, + }); + + expect(await backend.getCredentials("databricks")).toBeNull(); + expect(await backend.getCredentialStatus("databricks")).toMatchObject({ + configured: false, + authenticated: false, + readiness: "unauthenticated", + source: "oauth-m2m", + origin: "environment", + }); + }); + + it("uses the same complete M2M source for material and status", async () => { + const backend = createStoreBackend({ + store, + resolveAuthMethod, + oauthConfigForProvider: (_providerId, source) => + source?.type === "oauth-m2m" + ? { + grantType: "client-credentials", + clientId: source.clientId, + clientSecret: source.clientSecret, + tokenEndpoint: `${source.workspaceHost}/token`, + credentialBaseUrl: source.workspaceHost, + cacheKey: source.clientId, + } + : undefined, + env: { + DATABRICKS_AUTH_TYPE: "oauth-m2m", + DATABRICKS_TOKEN: "must-not-be-used", + DATABRICKS_HOST: "https://workspace.test", + DATABRICKS_CLIENT_ID: "client", + DATABRICKS_CLIENT_SECRET: "secret", + }, + }); + + expect(await backend.getCredentials("databricks")).toBeNull(); + expect(await backend.getCredentialStatus("databricks")).toMatchObject({ + configured: true, + authenticated: true, + source: "oauth-m2m", + origin: "environment", + }); + }); + }); + describe("onDidChangeCredentials", () => { it("is a no-op when no watched providers are given", () => { const backend = createStoreBackend({ store, resolveAuthMethod, env: {} }); diff --git a/packages/ai-credentials/src/store-backend/providerEnvMappings.ts b/packages/ai-credentials/src/store-backend/providerEnvMappings.ts index e5bb478..d81edb4 100644 --- a/packages/ai-credentials/src/store-backend/providerEnvMappings.ts +++ b/packages/ai-credentials/src/store-backend/providerEnvMappings.ts @@ -77,4 +77,7 @@ export const PROVIDER_ENV_MAPPINGS: Record = { deepseek: { apiKey: "DEEPSEEK_API_KEY", }, + databricks: { + apiKey: "DATABRICKS_TOKEN", + }, }; diff --git a/packages/ai-credentials/src/types/__tests__/credential-shaping.test.ts b/packages/ai-credentials/src/types/__tests__/credential-shaping.test.ts index c4c3c67..86e85b4 100644 --- a/packages/ai-credentials/src/types/__tests__/credential-shaping.test.ts +++ b/packages/ai-credentials/src/types/__tests__/credential-shaping.test.ts @@ -29,6 +29,7 @@ function fakeConfig(snowflake?: { host?: string; account?: string }): Credential getCustomHeaders: () => undefined, getAwsRegion: () => undefined, getSnowflake: () => snowflake, + getDatabricks: () => undefined, }; } @@ -39,6 +40,7 @@ function config(overrides: Partial = {}): CredentialConfig { getCustomHeaders: () => undefined, getAwsRegion: () => undefined, getSnowflake: () => undefined, + getDatabricks: () => undefined, ...overrides, }; } diff --git a/packages/ai-credentials/src/types/credential-shaping.ts b/packages/ai-credentials/src/types/credential-shaping.ts index 81f6237..d10a3e7 100644 --- a/packages/ai-credentials/src/types/credential-shaping.ts +++ b/packages/ai-credentials/src/types/credential-shaping.ts @@ -20,7 +20,11 @@ import type { ProviderCredentials } from "./credentials"; import type { Logger } from "./logger"; -import { buildSnowflakeCortexUrl, buildSnowflakeCortexUrlFromHost } from "./utils"; +import { + buildSnowflakeCortexUrl, + buildSnowflakeCortexUrlFromHost, + normalizeDatabricksHost, +} from "./utils"; /** * Maps a provider to its auth extension registration and credential type. @@ -59,6 +63,8 @@ export interface CredentialConfig { getAwsRegion(): string | undefined; /** Snowflake host/account (`authentication.snowflake.credentials`, env on the bridge side). */ getSnowflake(): { host?: string; account?: string } | undefined; + /** Databricks workspace host (`authentication.databricks.credentials`, env on the bridge side). */ + getDatabricks(): { host?: string } | undefined; } /** @@ -137,6 +143,13 @@ export function shapeCredentials( } else if (snowflake?.account) { baseUrl = buildSnowflakeCortexUrl(snowflake.account); } + } else if (mapping.authProviderId === "databricks") { + // Databricks workspace host, with env fallback for managed environments + // (e.g. Posit Workbench injecting DATABRICKS_HOST into sessions). + const databricks = config.getDatabricks(); + if (databricks?.host) { + baseUrl = normalizeDatabricksHost(databricks.host); + } } else { baseUrl = config.getBaseUrl(configKey) || undefined; } diff --git a/packages/ai-credentials/src/types/index.ts b/packages/ai-credentials/src/types/index.ts index 2961dcb..4fa0bc7 100644 --- a/packages/ai-credentials/src/types/index.ts +++ b/packages/ai-credentials/src/types/index.ts @@ -26,7 +26,11 @@ export type { AuthProviderMapping, CredentialConfig } from "./credential-shaping export type { Logger } from "./logger"; -export { buildSnowflakeCortexUrl, buildSnowflakeCortexUrlFromHost } from "./utils"; +export { + buildSnowflakeCortexUrl, + buildSnowflakeCortexUrlFromHost, + normalizeDatabricksHost, +} from "./utils"; // OAuth protocol/runtime types (moved from @assistant/core platform.ts) export type { DeviceAuthInfo, TokenData } from "./oauth-types"; diff --git a/packages/ai-credentials/src/types/utils.ts b/packages/ai-credentials/src/types/utils.ts index 73de041..429c46e 100644 --- a/packages/ai-credentials/src/types/utils.ts +++ b/packages/ai-credentials/src/types/utils.ts @@ -3,7 +3,7 @@ *--------------------------------------------------------------------------------------------*/ /** - * Snowflake Cortex URL construction helpers. + * Provider URL construction helpers (Snowflake Cortex, Databricks). * * Pure functions with no platform dependencies — safe for browser/renderer. */ @@ -32,3 +32,30 @@ export function buildSnowflakeCortexUrlFromHost(host: string): string { export function buildSnowflakeCortexUrl(account: string): string { return buildSnowflakeCortexUrlFromHost(`${account}.snowflakecomputing.com`); } + +// --------------------------------------------------------------------------- +// Databricks +// --------------------------------------------------------------------------- + +/** + * Normalize a Databricks workspace host to a bare `https://` origin. + * + * Accepts values with or without a scheme and with trailing slashes + * (users paste hosts in all of these shapes from the Databricks UI). + * + * @param raw - Workspace host (e.g. "adb-123.4.azuredatabricks.net/") + * @returns Normalized host (e.g. "https://adb-123.4.azuredatabricks.net") + */ +export function normalizeDatabricksHost(raw: string): string { + let value = raw.trim(); + if (!value) throw new Error("Databricks workspace URL is required"); + if (!/^[a-zA-Z][a-zA-Z0-9+.-]*:\/\//.test(value)) value = `https://${value}`; + const url = new URL(value); + if (url.protocol !== "https:") { + throw new Error("Databricks workspace URL must use HTTPS"); + } + if (url.username || url.password) { + throw new Error("Databricks workspace URL cannot contain credentials"); + } + return url.origin; +} diff --git a/packages/ai-provider-bridge/src/__tests__/credential-shaping.test.ts b/packages/ai-provider-bridge/src/__tests__/credential-shaping.test.ts index 403edf2..a0f427a 100644 --- a/packages/ai-provider-bridge/src/__tests__/credential-shaping.test.ts +++ b/packages/ai-provider-bridge/src/__tests__/credential-shaping.test.ts @@ -24,6 +24,7 @@ function fakeConfig( customHeaders?: Record>; awsRegion?: string; snowflake?: { host?: string; account?: string }; + databricks?: { host?: string }; } = {}, ): CredentialConfig { return { @@ -31,6 +32,7 @@ function fakeConfig( getCustomHeaders: (configKey) => overrides.customHeaders?.[configKey], getAwsRegion: () => overrides.awsRegion, getSnowflake: () => overrides.snowflake, + getDatabricks: () => overrides.databricks, }; } @@ -40,6 +42,7 @@ const OAUTH = { authProviderId: "posit-ai", credentialType: "oauth" } as const; const AWS = { authProviderId: "amazon-bedrock", credentialType: "aws-credentials" } as const; const GCP = { authProviderId: "google-cloud", credentialType: "google-cloud" } as const; const SNOWFLAKE = { authProviderId: "snowflake-cortex", credentialType: "apikey" } as const; +const DATABRICKS = { authProviderId: "databricks", credentialType: "apikey" } as const; describe("shapeCredentials", () => { // --- oauth --- @@ -113,6 +116,30 @@ describe("shapeCredentials", () => { }); }); + // --- databricks (apikey with normalized host URL) --- + + it("normalizes a databricks host from config to a bare https:// origin", () => { + const config = fakeConfig({ databricks: { host: "https://adb-123.4.azuredatabricks.net" } }); + expect(shapeCredentials(DATABRICKS, "dapi-tok", config)).toMatchObject({ + type: "apikey", + apiKey: "dapi-tok", + baseUrl: "https://adb-123.4.azuredatabricks.net", + }); + }); + + it("adds https:// scheme to a scheme-less databricks host and strips trailing slash", () => { + const config = fakeConfig({ databricks: { host: "my-workspace.cloud.databricks.com/" } }); + expect(shapeCredentials(DATABRICKS, "dapi-tok", config)).toMatchObject({ + baseUrl: "https://my-workspace.cloud.databricks.com", + }); + }); + + it("returns undefined baseUrl for databricks when no host is configured", () => { + expect(shapeCredentials(DATABRICKS, "dapi-tok", fakeConfig())).toMatchObject({ + baseUrl: undefined, + }); + }); + // --- aws-credentials --- it("shapes aws credentials, taking region from config", () => { diff --git a/packages/ai-provider-bridge/src/__tests__/register-all-providers.test.ts b/packages/ai-provider-bridge/src/__tests__/register-all-providers.test.ts index d2cab34..7456384 100644 --- a/packages/ai-provider-bridge/src/__tests__/register-all-providers.test.ts +++ b/packages/ai-provider-bridge/src/__tests__/register-all-providers.test.ts @@ -27,10 +27,12 @@ vi.mock("../providers/snowflake-cortex-provider", () => ({ registerSnowflakeCortexProvider: vi.fn(), })); vi.mock("../providers/deepseek-provider", () => ({ registerDeepSeekProvider: vi.fn() })); +vi.mock("../providers/databricks-provider", () => ({ registerDatabricksProvider: vi.fn() })); import { registerAnthropicProvider } from "../providers/anthropic-provider"; import { registerBedrockProvider } from "../providers/bedrock-provider"; import { registerCopilotProvider } from "../providers/copilot-provider"; +import { registerDatabricksProvider } from "../providers/databricks-provider"; import { registerDeepSeekProvider } from "../providers/deepseek-provider"; import { registerFoundryProvider } from "../providers/foundry-provider"; import { registerGeminiProvider } from "../providers/gemini-provider"; @@ -70,6 +72,7 @@ const allRegisterFns = [ registerFoundryProvider, registerSnowflakeCortexProvider, registerDeepSeekProvider, + registerDatabricksProvider, ]; describe("registerAllProviders (internal)", () => { @@ -89,7 +92,7 @@ describe("registerAllProviders (internal)", () => { expect(new Set(ids)).toEqual(new Set(PROVIDER_IDS)); }); - it("registers all 14 providers when allowedProviders is omitted", () => { + it("registers every provider when allowedProviders is omitted", () => { registerAllProviders(registry, mockLogger, { positAiBaseUrl: BASE_URL, userAgent: USER_AGENT }); // Tie the count to PROVIDER_IDS (the source of truth) so a provider added there diff --git a/packages/ai-provider-bridge/src/index.ts b/packages/ai-provider-bridge/src/index.ts index 44828ad..2232c97 100644 --- a/packages/ai-provider-bridge/src/index.ts +++ b/packages/ai-provider-bridge/src/index.ts @@ -102,6 +102,7 @@ export { normalizeBaseUrlForProvider } from "./base-url"; // Small utilities export { isThinkingEnabled } from "./utils"; export { buildSnowflakeCortexUrl } from "./utils"; +export { normalizeDatabricksHost } from "./utils"; export { isAgreementRequiredBody } from "./utils"; export { joinPath } from "./utils"; diff --git a/packages/ai-provider-bridge/src/model-capabilities/databricks-helpers.ts b/packages/ai-provider-bridge/src/model-capabilities/databricks-helpers.ts new file mode 100644 index 0000000..ded1697 --- /dev/null +++ b/packages/ai-provider-bridge/src/model-capabilities/databricks-helpers.ts @@ -0,0 +1,79 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (C) 2026 Posit Software, PBC. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +/** + * Capability inference for Databricks serving endpoints. + * + * The serving-endpoints list API exposes no token-limit or modality metadata, + * so capabilities are inferred from the underlying model identity (foundation + * model name, external model name, or endpoint name). Unknown identities get + * no override and fall back to the provider's conservative defaults. + */ + +import { getAnthropicModelCapabilities, getOpenAIModelCapabilities } from "ai-config"; + +import type { ModelInfo } from "../types"; + +/** + * Image MIME types accepted through the Databricks OpenAI-compatible surface. + * PDF input (which the upstream Anthropic/OpenAI helpers include) is not + * reliably supported across Databricks serving endpoints, so it is excluded. + */ +const DATABRICKS_IMAGE_MEDIA_TYPES = ["image/png", "image/jpeg", "image/gif", "image/webp"]; + +/** + * Strip Databricks-specific prefixes to get the bare upstream model identity. + * + * Handles: + * - Pay-per-token foundation models: `databricks-claude-sonnet-4-5` + * - Unity Catalog system models: `system.ai.claude-sonnet-4-5` + */ +function normalizeDatabricksModelId(modelId: string): string { + return modelId.replace(/^system\.ai\./, "").replace(/^databricks-/, ""); +} + +/** + * Infer model capabilities for a Databricks serving endpoint from the + * underlying model identity. + * + * Thinking-effort levels are deliberately dropped: reasoning-effort support is + * inconsistent across the Databricks OpenAI-compatible surface, so v1 does not + * offer thinking controls for Databricks models. + * + * @param modelIdentity - Best-known model identity (foundation model name, + * external model name, entity name, or endpoint name) + * @returns A partial `ModelInfo` override, or `undefined` for unrecognized models. + */ +export function getDatabricksModelCapabilities( + modelIdentity: string, +): Partial | undefined { + const normalized = normalizeDatabricksModelId(modelIdentity); + + const claudeCapabilities = getAnthropicModelCapabilities(normalized); + if (claudeCapabilities) { + const { thinkingEffortLevels: _dropped, ...capabilities } = claudeCapabilities; + return { + ...capabilities, + supportsImages: true, + supportedInputMediaTypes: DATABRICKS_IMAGE_MEDIA_TYPES, + }; + } + + const openaiCapabilities = getOpenAIModelCapabilities(normalized); + if (openaiCapabilities) { + const { + thinkingEffortLevels: _dropped, + supportedInputMediaTypes: _droppedMediaTypes, + ...capabilities + } = openaiCapabilities; + return { + ...capabilities, + ...(capabilities.supportsImages + ? { supportedInputMediaTypes: DATABRICKS_IMAGE_MEDIA_TYPES } + : {}), + }; + } + + return undefined; +} diff --git a/packages/ai-provider-bridge/src/provider-map.ts b/packages/ai-provider-bridge/src/provider-map.ts index 931a343..74ceb57 100644 --- a/packages/ai-provider-bridge/src/provider-map.ts +++ b/packages/ai-provider-bridge/src/provider-map.ts @@ -53,6 +53,10 @@ export const PROVIDER_MAP: Partial> = { credentialType: "apikey", }, deepseek: { authProviderId: "deepseek-api", scopes: [], credentialType: "apikey" }, + // Bearer-token provider: session.accessToken is a PAT or OAuth access token + // (the auth extension decides which); the workspace host comes from the + // `authentication.databricks.credentials` setting. + databricks: { authProviderId: "databricks", scopes: [], credentialType: "apikey" }, "google-vertex": { authProviderId: "google-cloud", scopes: [], diff --git a/packages/ai-provider-bridge/src/providers.ts b/packages/ai-provider-bridge/src/providers.ts index b95eb4b..eec8c60 100644 --- a/packages/ai-provider-bridge/src/providers.ts +++ b/packages/ai-provider-bridge/src/providers.ts @@ -11,6 +11,7 @@ // Provider registration functions export { registerAnthropicProvider } from "./providers/anthropic-provider"; export { registerCopilotProvider } from "./providers/copilot-provider"; +export { registerDatabricksProvider } from "./providers/databricks-provider"; export { registerDeepSeekProvider } from "./providers/deepseek-provider"; export { registerBedrockProvider } from "./providers/bedrock-provider"; export type { BedrockProviderCallbacks } from "./providers/bedrock-provider"; diff --git a/packages/ai-provider-bridge/src/providers/__tests__/databricks-provider.test.ts b/packages/ai-provider-bridge/src/providers/__tests__/databricks-provider.test.ts new file mode 100644 index 0000000..d18ee48 --- /dev/null +++ b/packages/ai-provider-bridge/src/providers/__tests__/databricks-provider.test.ts @@ -0,0 +1,518 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (C) 2026 Posit Software, PBC. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +import { OpenAIClient } from "../../model-clients/OpenAIClient"; +import type { Logger } from "../../types"; +import { + clearDatabricksGatewayModeCache, + parseFoundationModelsResponse, + parseServingEndpointsResponse, + registerDatabricksProvider, + rewriteServingUrlToGateway, +} from "../databricks-provider"; +import { ProviderRegistry } from "../ProviderRegistry"; + +const mockLogger: Logger = { + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + debug: vi.fn(), + trace: vi.fn(), +}; + +const HOST = "https://adb-123.4.azuredatabricks.net"; +const PROBE_URL = `${HOST}/api/ai-gateway/v2/endpoints?page_size=1`; +const SERVING_LIST_URL = `${HOST}/api/2.0/serving-endpoints`; +const FOUNDATION_LIST_URL = `${HOST}/api/2.0/serving-endpoints:foundation-models`; + +const CREDENTIALS = { + type: "apikey" as const, + apiKey: "dapi-test-token", + baseUrl: HOST, +}; + +/** Serving-endpoints list fixture covering every filter branch. */ +const SERVING_ENDPOINTS_FIXTURE = { + endpoints: [ + // FMAPI pay-per-token chat endpoint (foundation model) + { + name: "databricks-claude-sonnet-4-5", + task: "llm/v1/chat", + state: { ready: "READY", config_update: "NOT_UPDATING" }, + config: { + served_entities: [ + { + foundation_model: { + name: "databricks-claude-sonnet-4-5", + display_name: "Claude Sonnet 4.5", + }, + }, + ], + }, + }, + // External-model chat endpoint (task on the served entity, not top level) + { + name: "my-gpt-4o-gateway", + state: { ready: "READY" }, + config: { + served_entities: [ + { external_model: { provider: "openai", name: "gpt-4o", task: "llm/v1/chat" } }, + ], + }, + }, + // Custom chat endpoint with an unrecognized underlying model + { + name: "my-custom-chat-model", + task: "llm/v1/chat", + state: { ready: "READY" }, + config: { served_entities: [{ entity_name: "main.models.my_custom_model" }] }, + }, + // Embeddings endpoint — excluded (not chat) + { + name: "databricks-gte-large-en", + task: "llm/v1/embeddings", + state: { ready: "READY" }, + }, + // Completions-only endpoint — excluded (not chat) + { + name: "legacy-completions", + task: "llm/v1/completions", + state: { ready: "READY" }, + }, + // Chat endpoint that is not ready — excluded + { + name: "provisioning-chat", + task: "llm/v1/chat", + state: { ready: "NOT_READY" }, + }, + // Custom endpoint with no task at all — excluded + { + name: "feature-serving-endpoint", + state: { ready: "READY" }, + }, + // Route-optimized requires endpoint-scoped authorization_details — excluded. + { + name: "route-optimized-chat", + route_optimized: true, + task: "llm/v1/chat", + state: { ready: "READY" }, + }, + ], +}; + +/** Foundation-models list fixture (gateway discovery) covering the api_types filter. */ +const FOUNDATION_MODELS_FIXTURE = { + endpoints: [ + // Chat-capable foundation model with gateway v2 support + { + name: "databricks-claude-opus-4-8", + config: { + served_entities: [ + { + foundation_model: { + name: "databricks-claude-opus-4-8", + display_name: "Claude Opus 4.8", + api_types: [ + "mlflow/v1/chat/completions", + "anthropic/v1/messages", + "cursor/v1/chat/completions", + ], + ai_gateway_v2_supported: true, + }, + }, + ], + }, + }, + // Chat-capable open model + { + name: "databricks-llama-4-maverick", + config: { + served_entities: [ + { + foundation_model: { + name: "databricks-llama-4-maverick", + display_name: "Llama 4 Maverick", + api_types: ["mlflow/v1/chat/completions"], + ai_gateway_v2_supported: true, + }, + }, + ], + }, + }, + // Embeddings model — excluded (no chat api_type) + { + name: "databricks-gte-large-en", + config: { + served_entities: [ + { + foundation_model: { + name: "databricks-gte-large-en", + api_types: ["mlflow/v1/embeddings"], + ai_gateway_v2_supported: true, + }, + }, + ], + }, + }, + // Chat api_type but no gateway v2 support — excluded + { + name: "legacy-v1-only-chat", + config: { + served_entities: [ + { + foundation_model: { + name: "legacy-v1-only-chat", + api_types: ["mlflow/v1/chat/completions"], + ai_gateway_v2_supported: false, + }, + }, + ], + }, + }, + { + name: "route-optimized-foundation", + route_optimized: true, + config: { + served_entities: [ + { + foundation_model: { + name: "route-optimized-foundation", + api_types: ["mlflow/v1/chat/completions"], + ai_gateway_v2_supported: true, + }, + }, + ], + }, + }, + ], +}; + +function jsonResponse(body: unknown, status = 200): Response { + return new Response(JSON.stringify(body), { + status, + headers: { "Content-Type": "application/json" }, + }); +} + +/** + * Stub global fetch with a URL router. `probeStatus` controls the gateway + * availability probe; list URLs serve their fixtures. + */ +function stubRoutedFetch(options: { + probeStatus: number; + servingStatus?: number; + foundationStatus?: number; +}): ReturnType { + const fetchMock = vi.fn(async (input: string | URL | Request) => { + const url = typeof input === "string" || input instanceof URL ? input.toString() : input.url; + if (url === PROBE_URL) { + return jsonResponse({ endpoints: [] }, options.probeStatus); + } + if (url === SERVING_LIST_URL) { + return jsonResponse(SERVING_ENDPOINTS_FIXTURE, options.servingStatus ?? 200); + } + if (url === FOUNDATION_LIST_URL) { + return jsonResponse(FOUNDATION_MODELS_FIXTURE, options.foundationStatus ?? 200); + } + return jsonResponse({ message: `unexpected URL: ${url}` }, 500); + }); + vi.stubGlobal("fetch", fetchMock); + return fetchMock; +} + +function listUrlsCalled(fetchMock: ReturnType): string[] { + return fetchMock.mock.calls.map((call) => { + const input = call[0] as string | URL | Request; + return typeof input === "string" || input instanceof URL ? input.toString() : input.url; + }); +} + +describe("parseServingEndpointsResponse", () => { + it("keeps only READY chat-capable endpoints", () => { + const models = parseServingEndpointsResponse(SERVING_ENDPOINTS_FIXTURE); + + expect(models.map((m) => m.id)).toEqual([ + "databricks-claude-sonnet-4-5", + "my-gpt-4o-gateway", + "my-custom-chat-model", + ]); + }); + + it("maps endpoint name to model id and foundation display name to model name", () => { + const models = parseServingEndpointsResponse(SERVING_ENDPOINTS_FIXTURE); + const claude = models.find((m) => m.id === "databricks-claude-sonnet-4-5"); + + expect(claude).toMatchObject({ + id: "databricks-claude-sonnet-4-5", + name: "Claude Sonnet 4.5", + providerId: "databricks", + vendor: "databricks", + protocol: "openai", + }); + }); + + it("infers Claude capabilities from the foundation model name", () => { + const models = parseServingEndpointsResponse(SERVING_ENDPOINTS_FIXTURE); + const claude = models.find((m) => m.id === "databricks-claude-sonnet-4-5"); + + expect(claude).toMatchObject({ + family: "claude-4.5", + maxContextLength: 200_000, + supportsImages: true, + supportsToolResultImages: true, + }); + // Thinking controls are not offered for Databricks in v1. + expect(claude?.thinkingEffortLevels).toBeUndefined(); + }); + + it("infers OpenAI capabilities from the external model name", () => { + const models = parseServingEndpointsResponse(SERVING_ENDPOINTS_FIXTURE); + const gpt = models.find((m) => m.id === "my-gpt-4o-gateway"); + + expect(gpt).toMatchObject({ + family: "gpt-4o", + maxContextLength: 128_000, + supportsImages: true, + }); + expect(gpt?.thinkingEffortLevels).toBeUndefined(); + }); + + it("applies conservative defaults for unrecognized models", () => { + const models = parseServingEndpointsResponse(SERVING_ENDPOINTS_FIXTURE); + const custom = models.find((m) => m.id === "my-custom-chat-model"); + + expect(custom).toMatchObject({ + name: "my-custom-chat-model", + supportsTools: true, + supportsImages: false, + maxContextLength: 128_000, + maxOutputTokens: 16_384, + }); + }); + + it("returns an empty list for a malformed response", () => { + expect(parseServingEndpointsResponse({})).toEqual([]); + expect(parseServingEndpointsResponse({ endpoints: [] })).toEqual([]); + }); +}); + +describe("parseFoundationModelsResponse", () => { + it("keeps only gateway-v2 chat-capable models", () => { + const models = parseFoundationModelsResponse(FOUNDATION_MODELS_FIXTURE); + + expect(models.map((m) => m.id)).toEqual([ + "databricks-claude-opus-4-8", + "databricks-llama-4-maverick", + ]); + }); + + it("maps display names and infers capabilities from the foundation model name", () => { + const models = parseFoundationModelsResponse(FOUNDATION_MODELS_FIXTURE); + const opus = models.find((m) => m.id === "databricks-claude-opus-4-8"); + + expect(opus).toMatchObject({ + name: "Claude Opus 4.8", + providerId: "databricks", + vendor: "databricks", + protocol: "openai", + family: "claude-4.8", + maxContextLength: 1_000_000, + supportsImages: true, + }); + expect(opus?.thinkingEffortLevels).toBeUndefined(); + }); + + it("returns an empty list for a malformed response", () => { + expect(parseFoundationModelsResponse({})).toEqual([]); + expect(parseFoundationModelsResponse({ endpoints: [] })).toEqual([]); + }); +}); + +describe("rewriteServingUrlToGateway", () => { + it("rewrites the serving chat path to the gateway base path", () => { + expect(rewriteServingUrlToGateway(`${HOST}/serving-endpoints/chat/completions`, HOST)).toBe( + `${HOST}/ai-gateway/mlflow/v1/chat/completions`, + ); + }); + + it("leaves non-serving URLs untouched", () => { + expect(rewriteServingUrlToGateway(`${HOST}/api/2.0/something`, HOST)).toBe( + `${HOST}/api/2.0/something`, + ); + expect(rewriteServingUrlToGateway("https://other.example.com/serving-endpoints/x", HOST)).toBe( + "https://other.example.com/serving-endpoints/x", + ); + }); +}); + +describe("registerDatabricksProvider model fetcher", () => { + let registry: ProviderRegistry; + + beforeEach(() => { + vi.clearAllMocks(); + clearDatabricksGatewayModeCache(); + registry = new ProviderRegistry(mockLogger); + registerDatabricksProvider(registry, mockLogger); + }); + + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it("uses serving-endpoints discovery when the gateway probe returns 404", async () => { + const fetchMock = stubRoutedFetch({ probeStatus: 404 }); + + const models = await registry.getModelsForProvider("databricks", CREDENTIALS); + + expect(models.map((m) => m.id)).toEqual([ + "databricks-claude-sonnet-4-5", + "my-gpt-4o-gateway", + "my-custom-chat-model", + ]); + expect(listUrlsCalled(fetchMock)).toEqual([PROBE_URL, SERVING_LIST_URL]); + }); + + it("uses foundation-models discovery when the gateway probe succeeds", async () => { + const fetchMock = stubRoutedFetch({ probeStatus: 200 }); + + const models = await registry.getModelsForProvider("databricks", CREDENTIALS); + + expect(models.map((m) => m.id)).toEqual([ + "databricks-claude-opus-4-8", + "databricks-llama-4-maverick", + ]); + expect(listUrlsCalled(fetchMock)).toEqual([PROBE_URL, FOUNDATION_LIST_URL]); + }); + + it("falls back to serving mode without caching when the probe returns 5xx", async () => { + const fetchMock = stubRoutedFetch({ probeStatus: 503 }); + + const models = await registry.getModelsForProvider("databricks", CREDENTIALS); + + expect(models.map((m) => m.id)).toContain("databricks-claude-sonnet-4-5"); + expect(listUrlsCalled(fetchMock)).toEqual([PROBE_URL, SERVING_LIST_URL]); + + // A fresh fetcher (new registry, shared module cache) must probe again — + // the transient failure was not cached as a definitive answer. + const secondRegistry = new ProviderRegistry(mockLogger); + registerDatabricksProvider(secondRegistry, mockLogger); + await secondRegistry.getModelsForProvider("databricks", CREDENTIALS); + expect(listUrlsCalled(fetchMock).filter((u) => u === PROBE_URL)).toHaveLength(2); + }); + + it("caches the definitive gateway probe result across fetcher instances", async () => { + const fetchMock = stubRoutedFetch({ probeStatus: 200 }); + + await registry.getModelsForProvider("databricks", CREDENTIALS); + + const secondRegistry = new ProviderRegistry(mockLogger); + registerDatabricksProvider(secondRegistry, mockLogger); + await secondRegistry.getModelsForProvider("databricks", CREDENTIALS); + + expect(listUrlsCalled(fetchMock).filter((u) => u === PROBE_URL)).toHaveLength(1); + }); + + it("sends a Bearer token and additive customHeaders on probe and discovery", async () => { + const fetchMock = stubRoutedFetch({ probeStatus: 200 }); + + await registry.getModelsForProvider("databricks", { + ...CREDENTIALS, + customHeaders: { "x-databricks-use-coding-agent-mode": "true" }, + }); + + for (const call of fetchMock.mock.calls) { + expect(call[1]).toEqual({ + headers: { + Authorization: "Bearer dapi-test-token", + "x-databricks-use-coding-agent-mode": "true", + }, + }); + } + }); + + it("normalizes a scheme-less workspace host", async () => { + const fetchMock = stubRoutedFetch({ probeStatus: 404 }); + + await registry.getModelsForProvider("databricks", { + ...CREDENTIALS, + baseUrl: "adb-123.4.azuredatabricks.net/", + }); + + expect(listUrlsCalled(fetchMock)).toEqual([PROBE_URL, SERVING_LIST_URL]); + }); + + it("returns empty list when the API key is missing", async () => { + const fetchMock = stubRoutedFetch({ probeStatus: 200 }); + + const models = await registry.getModelsForProvider("databricks", { + type: "apikey", + apiKey: "", + baseUrl: HOST, + }); + + expect(models).toEqual([]); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it("returns empty list when the workspace host is missing", async () => { + const fetchMock = stubRoutedFetch({ probeStatus: 200 }); + + const models = await registry.getModelsForProvider("databricks", { + type: "apikey", + apiKey: "dapi-test-token", + }); + + expect(models).toEqual([]); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it("returns empty list when the discovery call fails", async () => { + stubRoutedFetch({ probeStatus: 404, servingStatus: 401 }); + + const models = await registry.getModelsForProvider("databricks", CREDENTIALS); + + expect(models).toEqual([]); + }); +}); + +describe("registerDatabricksProvider client factory", () => { + let registry: ProviderRegistry; + + beforeEach(() => { + vi.clearAllMocks(); + clearDatabricksGatewayModeCache(); + registry = new ProviderRegistry(mockLogger); + registerDatabricksProvider(registry, mockLogger); + }); + + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it("creates an OpenAI-compatible client for apikey credentials", () => { + const client = registry.getClientForProvider("databricks", CREDENTIALS); + + expect(client).toBeInstanceOf(OpenAIClient); + }); + + it("throws for non-apikey credentials", () => { + expect(() => + registry.getClientForProvider("databricks", { + type: "oauth", + accessToken: "some-token", + }), + ).toThrow(/requires API key credentials/); + }); + + it("throws when the workspace host is missing", () => { + expect(() => + registry.getClientForProvider("databricks", { + type: "apikey", + apiKey: "dapi-test-token", + }), + ).toThrow(/workspace host/); + }); +}); diff --git a/packages/ai-provider-bridge/src/providers/databricks-provider.ts b/packages/ai-provider-bridge/src/providers/databricks-provider.ts new file mode 100644 index 0000000..a39b281 --- /dev/null +++ b/packages/ai-provider-bridge/src/providers/databricks-provider.ts @@ -0,0 +1,362 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (C) 2026 Posit Software, PBC. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +/** + * Databricks provider + * + * Routes through Unity AI Gateway when the workspace has it enabled (Beta, + * account-level preview), falling back to classic Model Serving otherwise. + * Gateway routing gives requests centralized governance: usage tracking, + * inference tables, rate limits, and guardrails. + * + * Mode detection probes `GET {host}/api/ai-gateway/v2/endpoints` once per + * workspace host (cached; cleared with the model cache so credential or + * preview changes are picked up). + * + * - Gateway mode: models come from `GET {host}/api/2.0/serving-endpoints:foundation-models`, + * filtered to entities whose `foundation_model.api_types` include the + * MLflow chat-completions API. Chat goes to `{host}/ai-gateway/mlflow/v1/chat/completions`. + * - Serving mode: models come from `GET {host}/api/2.0/serving-endpoints`, + * filtered to READY chat-capable endpoints (task `llm/v1/chat`, including + * external-model entities). Chat goes to `{host}/serving-endpoints/chat/completions`. + * + * Both surfaces are OpenAI-compatible with the endpoint name as the `model` + * body parameter, so a single client serves both; only the base URL differs. + * + * Credentials are bearer-token `apikey` credentials: `apiKey` is a personal + * access token or an OAuth access token (the host application decides which), + * and `baseUrl` is the workspace host. + */ + +import { additiveHeaderRecord } from "../custom-headers"; +import { getDatabricksModelCapabilities } from "../model-capabilities/databricks-helpers"; +import { createOpenAICompatibleFetch } from "../model-clients/openai-compat-fetch"; +import { OpenAIClient } from "../model-clients/OpenAIClient"; +import type { ApiKeyCredentials, Logger, ModelInfo, ProviderCredentials } from "../types"; +import { normalizeDatabricksHost } from "../utils"; +import type { ClearableModelFetcher } from "./cached-model-fetcher"; +import type { ProviderRegistry } from "./ProviderRegistry"; + +const CACHE_TTL = 60 * 60 * 1000; // 60 minutes, matching createCachedModelFetcher + +/** Conservative defaults for endpoints whose underlying model is unrecognized. */ +const DATABRICKS_DEFAULTS = { + vendor: "databricks" as const, + protocol: "openai" as const, + supportsTools: true, + supportsImages: false, + supportsToolResultImages: false, + supportsWebSearch: false, + maxInputTokens: 128_000, + maxOutputTokens: 16_384, + maxContextLength: 128_000, +} satisfies Partial; + +/** Serving-endpoint task indicating an OpenAI-style chat interface. */ +const CHAT_TASK = "llm/v1/chat"; + +/** Unity AI Gateway api_type indicating the MLflow chat-completions API. */ +const GATEWAY_CHAT_API_TYPE = "mlflow/v1/chat/completions"; + +/** Path suffixes for the two chat surfaces (appended to the workspace host). */ +const SERVING_CHAT_BASE_PATH = "/serving-endpoints"; +const GATEWAY_CHAT_BASE_PATH = "/ai-gateway/mlflow/v1"; + +// --------------------------------------------------------------------------- +// Gateway availability probe (cached per workspace host) +// --------------------------------------------------------------------------- + +/** + * Per-host Unity AI Gateway availability. Shared between the model fetcher + * and the chat client so both route consistently. Only definitive probe + * results are cached; transient failures fall back to serving mode for that + * attempt without poisoning the cache. + */ +const gatewayModeCache = new Map(); + +/** Clear cached gateway availability (exported for tests; also cleared with the model cache). */ +export function clearDatabricksGatewayModeCache(): void { + gatewayModeCache.clear(); +} + +/** + * Whether the workspace has Unity AI Gateway enabled. + * 200 and 404/403 are definitive and cached; anything else (5xx, network + * errors) defaults to serving mode for this attempt only. + */ +async function resolveGatewayMode( + host: string, + headers: Record, + logger: Logger, +): Promise { + const cached = gatewayModeCache.get(host); + if (cached !== undefined) { + return cached; + } + + try { + const response = await fetch(`${host}/api/ai-gateway/v2/endpoints?page_size=1`, { headers }); + if (response.ok) { + logger.debug(`[databricks] Unity AI Gateway available; routing via gateway`); + gatewayModeCache.set(host, true); + return true; + } + if (response.status === 404 || response.status === 403) { + logger.debug( + `[databricks] Unity AI Gateway unavailable (${response.status}); routing via model serving`, + ); + gatewayModeCache.set(host, false); + return false; + } + logger.warn( + `[databricks] Gateway probe returned ${response.status}; using model serving for this attempt`, + ); + return false; + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + logger.warn(`[databricks] Gateway probe failed: ${message}; using model serving`); + return false; + } +} + +// --------------------------------------------------------------------------- +// Response parsing +// --------------------------------------------------------------------------- + +interface FoundationModel { + name?: string; + display_name?: string; + api_types?: string[]; + ai_gateway_v2_supported?: boolean; +} + +interface ServedEntity { + entity_name?: string; + foundation_model?: FoundationModel; + external_model?: { name?: string; provider?: string; task?: string }; +} + +interface ServingEndpoint { + name?: string; + task?: string; + /** Requires endpoint-scoped OAuth authorization_details; excluded for now. */ + route_optimized?: boolean; + state?: { ready?: string }; + config?: { served_entities?: ServedEntity[] }; +} + +/** + * Whether a serving endpoint exposes a chat interface. Pay-per-token and + * provisioned-throughput endpoints carry the task at the top level; + * external-model endpoints carry it on each served entity. + */ +function isChatEndpoint(endpoint: ServingEndpoint): boolean { + if (endpoint.task === CHAT_TASK) { + return true; + } + const entities = endpoint.config?.served_entities ?? []; + return entities.some((entity) => entity.external_model?.task === CHAT_TASK); +} + +/** + * Best-known identity of the model behind an endpoint, used for capability + * inference. Falls back to the endpoint name (pay-per-token endpoint names + * like `databricks-claude-sonnet-4-5` identify the model directly). + */ +function resolveModelIdentity(endpoint: ServingEndpoint): string { + const entity = endpoint.config?.served_entities?.[0]; + return ( + entity?.foundation_model?.name ?? + entity?.external_model?.name ?? + entity?.entity_name ?? + endpoint.name ?? + "" + ); +} + +function toModelInfo(endpoint: ServingEndpoint): ModelInfo { + const displayName = endpoint.config?.served_entities?.[0]?.foundation_model?.display_name; + return { + id: endpoint.name ?? "", + name: displayName ?? endpoint.name ?? "", + providerId: "databricks", + ...DATABRICKS_DEFAULTS, + ...getDatabricksModelCapabilities(resolveModelIdentity(endpoint)), + }; +} + +/** + * Parse a serving-endpoints list response into chat-capable models. + * Exported for tests. + */ +export function parseServingEndpointsResponse(data: unknown): ModelInfo[] { + const endpoints = (data as { endpoints?: ServingEndpoint[] }).endpoints ?? []; + + const models: ModelInfo[] = []; + for (const endpoint of endpoints) { + if (!endpoint.name) continue; + if (endpoint.route_optimized === true) continue; + if (endpoint.state?.ready !== "READY") continue; + if (!isChatEndpoint(endpoint)) continue; + models.push(toModelInfo(endpoint)); + } + return models; +} + +/** + * Parse a foundation-models list response into gateway-chat-capable models. + * Exported for tests. + */ +export function parseFoundationModelsResponse(data: unknown): ModelInfo[] { + const endpoints = (data as { endpoints?: ServingEndpoint[] }).endpoints ?? []; + + const models: ModelInfo[] = []; + for (const endpoint of endpoints) { + if (!endpoint.name) continue; + if (endpoint.route_optimized === true) continue; + const entities = endpoint.config?.served_entities ?? []; + const gatewayChatCapable = entities.some( + (entity) => + entity.foundation_model?.ai_gateway_v2_supported === true && + (entity.foundation_model.api_types ?? []).includes(GATEWAY_CHAT_API_TYPE), + ); + if (!gatewayChatCapable) continue; + models.push(toModelInfo(endpoint)); + } + return models; +} + +// --------------------------------------------------------------------------- +// Registration +// --------------------------------------------------------------------------- + +async function fetchModelList(url: string, headers: Record): Promise { + const response = await fetch(url, { headers }); + if (!response.ok) { + throw new Error(`API returned ${response.status}`); + } + return response.json(); +} + +function createDatabricksModelFetcher(logger: Logger): ClearableModelFetcher { + let lastFetch = 0; + let cachedModels: ModelInfo[] | null = null; + + const fetcher: ClearableModelFetcher = async ( + credentials: ProviderCredentials, + ): Promise => { + const typed = credentials as ApiKeyCredentials; + if (!typed.apiKey || !typed.baseUrl?.trim()) { + logger.debug("[databricks] Missing apiKey or workspace host, returning no models"); + return []; + } + + const now = Date.now(); + if (cachedModels && now - lastFetch < CACHE_TTL) { + logger.debug("[databricks] Using cached models"); + return cachedModels; + } + + const host = normalizeDatabricksHost(typed.baseUrl); + const headers = additiveHeaderRecord( + { Authorization: `Bearer ${typed.apiKey}` }, + typed.customHeaders, + ); + + try { + const useGateway = await resolveGatewayMode(host, headers, logger); + const models = useGateway + ? parseFoundationModelsResponse( + await fetchModelList(`${host}/api/2.0/serving-endpoints:foundation-models`, headers), + ) + : parseServingEndpointsResponse( + await fetchModelList(`${host}/api/2.0/serving-endpoints`, headers), + ); + + lastFetch = now; + cachedModels = models; + logger.info( + `[databricks] Fetched ${models.length} chat models via ${useGateway ? "Unity AI Gateway" : "model serving"}`, + ); + return models; + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + logger.warn(`[databricks] Model fetch failed: ${message}`); + return cachedModels ?? []; + } + }; + + fetcher.clearCache = () => { + cachedModels = null; + lastFetch = 0; + clearDatabricksGatewayModeCache(); + }; + + return fetcher; +} + +/** + * Rewrite a model-serving chat URL to the Unity AI Gateway base path. + * Exported for tests. + */ +export function rewriteServingUrlToGateway(url: string, host: string): string { + const servingBase = `${host}${SERVING_CHAT_BASE_PATH}`; + if (url.startsWith(servingBase)) { + return `${host}${GATEWAY_CHAT_BASE_PATH}${url.slice(servingBase.length)}`; + } + return url; +} + +/** + * Wrap the OpenAI-compatible fetch so chat requests are re-routed to the + * Unity AI Gateway when the workspace supports it. The client is constructed + * against the model-serving base URL; this wrapper swaps the path prefix per + * request based on the cached gateway probe. + */ +function createDatabricksChatFetch( + host: string, + apiKey: string, + customHeaders: Record | undefined, + logger: Logger, +): typeof globalThis.fetch { + const compatFetch = createOpenAICompatibleFetch("Databricks", apiKey, customHeaders); + const probeHeaders = additiveHeaderRecord({ Authorization: `Bearer ${apiKey}` }, customHeaders); + + const fetchWithRouting = async ( + input: Parameters[0], + init?: Parameters[1], + ): Promise => { + const useGateway = await resolveGatewayMode(host, probeHeaders, logger); + if (!useGateway) { + return compatFetch(input, init); + } + if (typeof input === "string" || input instanceof URL) { + return compatFetch(rewriteServingUrlToGateway(input.toString(), host), init); + } + return compatFetch(new Request(rewriteServingUrlToGateway(input.url, host), input), init); + }; + + return fetchWithRouting as typeof globalThis.fetch; +} + +export function registerDatabricksProvider(registry: ProviderRegistry, logger: Logger): void { + registry.registerModelFetcher("databricks", createDatabricksModelFetcher(logger)); + + registry.registerClientFactory("databricks", (credentials) => { + if (credentials.type !== "apikey") { + throw new Error(`Databricks provider requires API key credentials, got: ${credentials.type}`); + } + if (!credentials.baseUrl?.trim()) { + throw new Error("Databricks provider requires a workspace host (baseUrl)"); + } + const host = normalizeDatabricksHost(credentials.baseUrl); + // customHeaders are injected by the custom fetch wrapper. + return new OpenAIClient( + credentials.apiKey, + `${host}${SERVING_CHAT_BASE_PATH}`, + "completions", + createDatabricksChatFetch(host, credentials.apiKey, credentials.customHeaders, logger), + ); + }); +} diff --git a/packages/ai-provider-bridge/src/register-all-providers.ts b/packages/ai-provider-bridge/src/register-all-providers.ts index 2bda72b..fc6e3ee 100644 --- a/packages/ai-provider-bridge/src/register-all-providers.ts +++ b/packages/ai-provider-bridge/src/register-all-providers.ts @@ -18,6 +18,7 @@ import { import { registerAnthropicProvider } from "./providers/anthropic-provider"; import { registerBedrockProvider } from "./providers/bedrock-provider"; import { registerCopilotProvider } from "./providers/copilot-provider"; +import { registerDatabricksProvider } from "./providers/databricks-provider"; import { registerDeepSeekProvider } from "./providers/deepseek-provider"; import { registerFoundryProvider } from "./providers/foundry-provider"; import { registerGeminiProvider } from "./providers/gemini-provider"; @@ -82,6 +83,7 @@ export const PROVIDER_REGISTRARS: readonly [ProviderId, ProviderRegistrar][] = [ ["ms-foundry", registerFoundryProvider], ["snowflake-cortex", registerSnowflakeCortexProvider], ["deepseek", registerDeepSeekProvider], + ["databricks", registerDatabricksProvider], ]; /** diff --git a/packages/ai-provider-bridge/src/types.ts b/packages/ai-provider-bridge/src/types.ts index e8d3caf..a7cc630 100644 --- a/packages/ai-provider-bridge/src/types.ts +++ b/packages/ai-provider-bridge/src/types.ts @@ -54,6 +54,7 @@ export const PROVIDER_IDS = [ "snowflake-cortex", "ms-foundry", "deepseek", + "databricks", ] as const; /** diff --git a/packages/ai-provider-bridge/src/utils.ts b/packages/ai-provider-bridge/src/utils.ts index 9a8fdd0..dee05d3 100644 --- a/packages/ai-provider-bridge/src/utils.ts +++ b/packages/ai-provider-bridge/src/utils.ts @@ -30,11 +30,15 @@ export function isClaudeModel(modelId: string): boolean { } // --------------------------------------------------------------------------- -// Snowflake +// Snowflake / Databricks // --------------------------------------------------------------------------- // Re-exported from ai-credentials/types (single source of truth) -export { buildSnowflakeCortexUrl, buildSnowflakeCortexUrlFromHost } from "ai-credentials/types"; +export { + buildSnowflakeCortexUrl, + buildSnowflakeCortexUrlFromHost, + normalizeDatabricksHost, +} from "ai-credentials/types"; // --------------------------------------------------------------------------- // Posit AI