diff --git a/packages/ai-credentials/src/positron/PositronBackend.ts b/packages/ai-credentials/src/positron/PositronBackend.ts index ad3bb3d..1419c91 100644 --- a/packages/ai-credentials/src/positron/PositronBackend.ts +++ b/packages/ai-credentials/src/positron/PositronBackend.ts @@ -49,8 +49,16 @@ export interface PositronBackend extends Backend { export interface CreatePositronBackendOptions { logger: Logger; - /** Provider-id → auth mapping (the bridge's PROVIDER_MAP). */ - providerMap: ProviderMap; + /** + * Provider-id → auth mapping (the bridge's `PROVIDER_MAP`, plus whatever the + * host adds for `providers.custom` entries). + * + * A getter, not a table: custom entry ids are user-chosen and come and go + * while the process runs, so a table read once at construction would silently + * stop resolving — and stop firing credential-change events — for anything + * added afterwards. + */ + providerMap: () => ProviderMap; /** CredentialConfig factory (the host injects its catalog-backed adapter). */ credentialConfigFactory: () => CredentialConfig; } @@ -99,8 +107,6 @@ async function tryCreateSession( export function createPositronBackend(options: CreatePositronBackendOptions): PositronBackend { const { logger, providerMap, credentialConfigFactory } = options; - const mappedProviderIds = Object.keys(providerMap).filter((id) => providerMap[id] !== undefined); - // Auth provider ids observed to time out waiting to register (i.e. not shipped // in this host build). Cached for the process lifetime so the multi-second // registration wait is NOT re-paid on every silent lookup — a conversation @@ -120,6 +126,18 @@ export function createPositronBackend(options: CreatePositronBackendOptions): Po // throws "No credentials available"). Correctness over shaving a one-time wait. const unregisteredAuthProviders = new Set(); + // Per-auth-provider count of session changes seen, i.e. of "it registered" + // signals. A lookup records the count it started under and only caches its + // verdict if the count still matches, because the registration timeout is + // several seconds long and the provider can register inside that window: the + // event's `delete` would land first and the stale `add` after it, leaving the + // provider permanently unresolvable with no further event to recover it. That + // is the ordering a user hits when a `providers.custom` entry is added or + // first loaded at the same time as its auth provider registers. + const registrationSignals = new Map(); + const signalsFor = (authProviderId: string): number => + registrationSignals.get(authProviderId) ?? 0; + async function trySilentSession( authProviderId: string, scopes: string[], @@ -128,20 +146,25 @@ export function createPositronBackend(options: CreatePositronBackendOptions): Po // for the full registration timeout again — skip the call entirely. if (unregisteredAuthProviders.has(authProviderId)) return undefined; + const signalsAtStart = signalsFor(authProviderId); try { const session = await vscode.authentication.getSession(authProviderId, scopes, { silent: true, }); return session ?? undefined; } catch (err) { - if (isProviderNotRegisteredError(err, authProviderId)) { - unregisteredAuthProviders.add(authProviderId); + if (!isProviderNotRegisteredError(err, authProviderId)) { + logger.debug( + `[ai-credentials/positron] Auth session unavailable for ${authProviderId}: ${err}`, + ); + } else if (signalsFor(authProviderId) !== signalsAtStart) { logger.trace( - `[ai-credentials/positron] Auth provider ${authProviderId} is not registered; skipping future silent lookups`, + `[ai-credentials/positron] Auth provider ${authProviderId} registered while this lookup was in flight; not caching the verdict`, ); } else { - logger.debug( - `[ai-credentials/positron] Auth session unavailable for ${authProviderId}: ${err}`, + unregisteredAuthProviders.add(authProviderId); + logger.trace( + `[ai-credentials/positron] Auth provider ${authProviderId} is not registered; skipping future silent lookups`, ); } return undefined; @@ -152,7 +175,7 @@ export function createPositronBackend(options: CreatePositronBackendOptions): Po providerId: string, prompt: boolean, ): Promise { - const mapping = providerMap[providerId]; + const mapping = providerMap()[providerId]; if (!mapping) return null; const { authProviderId, scopes, fallbackScopes } = mapping; @@ -171,22 +194,18 @@ export function createPositronBackend(options: CreatePositronBackendOptions): Po } if (!session) return null; - return shapeCredentials(mapping, session.accessToken, credentialConfigFactory(), logger); + return shapeCredentials( + providerId, + mapping, + session.accessToken, + credentialConfigFactory(), + logger, + ); } // --- Credential change events ------------------------------------------- const emitter = new vscode.EventEmitter(); - // Reverse map: auth provider id -> logical provider ids. - const authToLogical = new Map(); - for (const logicalId of mappedProviderIds) { - const mapping = providerMap[logicalId]; - if (!mapping) continue; - const list = authToLogical.get(mapping.authProviderId) ?? []; - list.push(logicalId); - authToLogical.set(mapping.authProviderId, list); - } - // The emitter fires ONLY on vscode auth session changes (login/logout). // // Connection-config changes (base URL, customHeaders, AWS region, Snowflake @@ -199,10 +218,18 @@ export function createPositronBackend(options: CreatePositronBackendOptions): Po // before the debounced rebuild lands. const sessionSub = vscode.authentication.onDidChangeSessions((e) => { // A session change means the provider is registered now: drop any stale - // "unregistered" verdict so silent lookups resume against it. + // "unregistered" verdict so silent lookups resume against it, and count the + // signal so a lookup still waiting on the registration timeout can't + // re-install the verdict when it finally rejects. + registrationSignals.set(e.provider.id, signalsFor(e.provider.id) + 1); unregisteredAuthProviders.delete(e.provider.id); - const logicalIds = authToLogical.get(e.provider.id); - if (logicalIds) emitter.fire(logicalIds); + // Reverse lookup per event rather than an index built at construction, so a + // custom entry added after this backend was created still notifies. + const current = providerMap(); + const logicalIds = Object.keys(current).filter( + (id) => current[id]?.authProviderId === e.provider.id, + ); + if (logicalIds.length > 0) emitter.fire(logicalIds); }); function onDidChangeCredentials(callback: (providerIds: string[]) => void): Disposable { diff --git a/packages/ai-credentials/src/positron/__tests__/PositronBackend.test.ts b/packages/ai-credentials/src/positron/__tests__/PositronBackend.test.ts index 962eb65..9b019e5 100644 --- a/packages/ai-credentials/src/positron/__tests__/PositronBackend.test.ts +++ b/packages/ai-credentials/src/positron/__tests__/PositronBackend.test.ts @@ -51,6 +51,7 @@ vi.mock("vscode", () => ({ const { createPositronBackend } = await import("../PositronBackend.js"); import type { AuthProviderMapping, CredentialConfig } from "../../types/index.js"; +import type { ProviderMap } from "../PositronBackend.js"; const PROVIDER_MAP: Record = { anthropic: { authProviderId: "anthropic-api", scopes: [], credentialType: "apikey" }, @@ -82,10 +83,13 @@ function testConfig(overrides: Partial = {}): CredentialConfig }; } -function makeBackend(configOverrides: Partial = {}) { +function makeBackend( + configOverrides: Partial = {}, + extraMappings?: () => ProviderMap, +) { return createPositronBackend({ logger, - providerMap: PROVIDER_MAP, + providerMap: () => ({ ...PROVIDER_MAP, ...extraMappings?.() }), credentialConfigFactory: () => testConfig(configOverrides), }); } @@ -203,7 +207,7 @@ describe("createPositronBackend", () => { mockGetSession.mockResolvedValue(makeSession("databricks-bearer-token")); const backend = makeBackend({ getDatabricks: () => ({ host: "https://adb-123.4.azuredatabricks.net" }), - getCustomHeaders: (configKey) => + getCustomHeaders: ({ configKey }) => configKey === "databricks" ? { "x-databricks-use-coding-agent-mode": "true" } : undefined, }); @@ -313,6 +317,31 @@ describe("createPositronBackend", () => { expect(mockGetSession).toHaveBeenCalledTimes(2); }); + it("does not cache the verdict when the provider registers mid-lookup", async () => { + // The registration timeout is several seconds long, so the provider can + // register while a lookup is still waiting on it. The session change lands + // first and the rejection after, so caching on rejection would leave the + // provider permanently unresolvable with no further event to recover it. + let rejectFirst!: (err: Error) => void; + mockGetSession.mockReturnValueOnce( + new Promise((_resolve, reject) => { + rejectFirst = reject; + }), + ); + const backend = makeBackend(); + const firstLookup = backend.getCredentials("databricks"); + + sessionChangeHook.callback?.({ provider: { id: "databricks" } }); + rejectFirst(NOT_REGISTERED); + expect(await firstLookup).toBeNull(); + + mockGetSession.mockResolvedValue(makeSession("databricks-bearer-token")); + expect(await backend.getCredentials("databricks")).toMatchObject({ + apiKey: "databricks-bearer-token", + }); + expect(mockGetSession).toHaveBeenCalledTimes(2); + }); + it("logs the registration timeout at trace, not debug", async () => { mockGetSession.mockRejectedValue(NOT_REGISTERED); const backend = makeBackend(); @@ -341,8 +370,9 @@ describe("createPositronBackend", () => { it("shapes baseUrl and customHeaders from the injected credential config", async () => { mockGetSession.mockResolvedValue(makeSession("sk-ant")); const backend = makeBackend({ - getBaseUrl: (configKey) => (configKey === "anthropic" ? "https://proxy.example" : undefined), - getCustomHeaders: (configKey) => + getBaseUrl: ({ configKey }) => + configKey === "anthropic" ? "https://proxy.example" : undefined, + getCustomHeaders: ({ configKey }) => configKey === "anthropic" ? { "x-tenancy": "team-42" } : undefined, }); @@ -353,4 +383,95 @@ describe("createPositronBackend", () => { customHeaders: { "x-tenancy": "team-42" }, }); }); + + describe("providers.custom entries", () => { + // The shape the Positron host composes for `providers.custom`: every + // entry shares ONE authentication provider and is identified by a single + // scope — the entry id. So `providerId !== authProviderId`, and the + // configKey derived from the shared authProviderId is the same for every + // entry; only the logical provider id tells two entries apart. + const CUSTOM_AUTH_PROVIDER_ID = "positron-custom-provider"; + const ACME = "Acme Gateway"; + const CONTOSO = "Contoso Gateway"; + + function customMappings(ids: readonly string[]): () => ProviderMap { + return () => + Object.fromEntries( + ids.map((id) => [ + id, + { + authProviderId: CUSTOM_AUTH_PROVIDER_ID, + scopes: [id], + credentialType: "apikey" as const, + }, + ]), + ); + } + + it("resolves each entry by scope, with shaping keyed on the entry id", async () => { + // Two entries behind one auth provider: the session lookup must pass + // the entry's scope, and config reads must key on the logical + // providerId — keyed on the shared authProviderId/configKey instead, + // both entries would resolve the same (or no) baseUrl. + mockGetSession.mockImplementation(async (_id: string, scopes: string[]) => + makeSession(`key-for-${scopes[0]}`), + ); + const backend = makeBackend( + { + getBaseUrl: ({ providerId }) => + providerId === ACME + ? "https://gw.acme.test" + : providerId === CONTOSO + ? "https://gw.contoso.test" + : undefined, + }, + customMappings([ACME, CONTOSO]), + ); + + await expect(backend.getCredentials(ACME)).resolves.toEqual({ + type: "apikey", + apiKey: "key-for-Acme Gateway", + baseUrl: "https://gw.acme.test", + customHeaders: undefined, + }); + await expect(backend.getCredentials(CONTOSO)).resolves.toEqual({ + type: "apikey", + apiKey: "key-for-Contoso Gateway", + baseUrl: "https://gw.contoso.test", + customHeaders: undefined, + }); + expect(mockGetSession).toHaveBeenNthCalledWith(1, CUSTOM_AUTH_PROVIDER_ID, [ACME], { + silent: true, + }); + expect(mockGetSession).toHaveBeenNthCalledWith(2, CUSTOM_AUTH_PROVIDER_ID, [CONTOSO], { + silent: true, + }); + }); + + it("fans a session change on the shared auth provider out to every entry", () => { + const backend = makeBackend({}, customMappings([ACME, CONTOSO])); + const seen: string[][] = []; + backend.onDidChangeCredentials((ids) => seen.push(ids)); + + sessionChangeHook.callback?.({ provider: { id: CUSTOM_AUTH_PROVIDER_ID } }); + + expect(seen).toEqual([[ACME, CONTOSO]]); + }); + + it("notifies a custom entry that appeared after the backend was built", () => { + // The map is read per lookup, not captured at construction, so an entry + // added later is still reachable. + const known: string[] = []; + const backend = makeBackend({}, customMappings(known)); + const seen: string[][] = []; + backend.onDidChangeCredentials((ids) => seen.push(ids)); + + sessionChangeHook.callback?.({ provider: { id: CUSTOM_AUTH_PROVIDER_ID } }); + expect(seen).toEqual([]); + + known.push("Late Entry"); + sessionChangeHook.callback?.({ provider: { id: CUSTOM_AUTH_PROVIDER_ID } }); + expect(seen).toEqual([["Late Entry"]]); + }); + }); }); diff --git a/packages/ai-credentials/src/store/__tests__/helpers/lock-holder.ts b/packages/ai-credentials/src/store/__tests__/helpers/lock-holder.ts index 7c55b00..d038ed1 100644 --- a/packages/ai-credentials/src/store/__tests__/helpers/lock-holder.ts +++ b/packages/ai-credentials/src/store/__tests__/helpers/lock-holder.ts @@ -42,6 +42,8 @@ async function main() { resolve(); return; } + // Pass sendHandle/options explicitly: the short send(message, callback) + // overload only exists in @types/node >= 22.19, and consumers can hoist older. process.send("lock-released", undefined, undefined, () => resolve()); }); process.disconnect(); 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 bfa5028..38c714e 100644 --- a/packages/ai-credentials/src/types/__tests__/credential-shaping.test.ts +++ b/packages/ai-credentials/src/types/__tests__/credential-shaping.test.ts @@ -48,7 +48,7 @@ function config(overrides: Partial = {}): CredentialConfig { describe("shapeCredentials — Snowflake host-over-account URL", () => { it("builds the URL from host, not account, when both are present", () => { const config = fakeConfig({ host: "h.snowflakecomputing.com", account: "org-acct" }); - expect(shapeCredentials(SNOWFLAKE, "tok", config)).toMatchObject({ + expect(shapeCredentials("snowflake-cortex", SNOWFLAKE, "tok", config)).toMatchObject({ type: "apikey", baseUrl: "https://h.snowflakecomputing.com/api/v2/cortex/v1", }); @@ -56,16 +56,41 @@ describe("shapeCredentials — Snowflake host-over-account URL", () => { it("falls back to account when only account is present", () => { const config = fakeConfig({ account: "org-acct" }); - expect(shapeCredentials(SNOWFLAKE, "tok", config)).toMatchObject({ + expect(shapeCredentials("snowflake-cortex", SNOWFLAKE, "tok", config)).toMatchObject({ baseUrl: "https://org-acct.snowflakecomputing.com/api/v2/cortex/v1", }); }); it("leaves the URL undefined when neither host nor account is present", () => { - expect(shapeCredentials(SNOWFLAKE, "tok", fakeConfig())).toMatchObject({ + expect(shapeCredentials("snowflake-cortex", SNOWFLAKE, "tok", fakeConfig())).toMatchObject({ baseUrl: undefined, }); }); + + // Standalone's Add-custom-provider form writes a flat `baseUrl` in custom-URL + // mode and structured host/account otherwise, into the same providers.json, + // so both shapes have to resolve. When both are present the flat form wins, + // matching the Node catalog paths (`conn.baseUrl ?? derive(conn)`), so one + // providers.json can't route different hosts to different endpoints. + it("falls back to structured fields when there is no flat baseUrl", () => { + const cfg = config({ + getBaseUrl: () => undefined, + getSnowflake: () => ({ host: "h.snowflakecomputing.com" }), + }); + expect(shapeCredentials("snowflake-cortex", SNOWFLAKE, "tok", cfg)).toMatchObject({ + baseUrl: "https://h.snowflakecomputing.com/api/v2/cortex/v1", + }); + }); + + it("prefers a flat baseUrl over structured fields when both are present", () => { + const cfg = config({ + getBaseUrl: () => "https://proxy.example.com/cortex/v1", + getSnowflake: () => ({ host: "h.snowflakecomputing.com" }), + }); + expect(shapeCredentials("snowflake-cortex", SNOWFLAKE, "tok", cfg)).toMatchObject({ + baseUrl: "https://proxy.example.com/cortex/v1", + }); + }); }); // Parity coverage ported from the removed ai-provider-bridge positron auth suite. @@ -79,7 +104,12 @@ describe("shapeCredentials — AWS credentials JSON", () => { it("parses the JSON token and applies the configured region", () => { expect( - shapeCredentials(AWS, awsToken, config({ getAws: () => ({ region: "eu-west-1" }) })), + shapeCredentials( + "bedrock", + AWS, + awsToken, + config({ getAws: () => ({ region: "eu-west-1" }) }), + ), ).toEqual({ type: "aws-credentials", region: "eu-west-1", @@ -90,20 +120,24 @@ describe("shapeCredentials — AWS credentials JSON", () => { }); it("defaults the region to us-east-1 when none is configured", () => { - expect(shapeCredentials(AWS, awsToken, config())).toMatchObject({ region: "us-east-1" }); + expect(shapeCredentials("bedrock", AWS, awsToken, config())).toMatchObject({ + region: "us-east-1", + }); }); it("returns null for a non-JSON token", () => { - expect(shapeCredentials(AWS, "not-json", config())).toBeNull(); + expect(shapeCredentials("bedrock", AWS, "not-json", config())).toBeNull(); }); it("returns null when accessKeyId or secretAccessKey is missing", () => { - expect(shapeCredentials(AWS, JSON.stringify({ accessKeyId: "AKIA" }), config())).toBeNull(); + expect( + shapeCredentials("bedrock", AWS, JSON.stringify({ accessKeyId: "AKIA" }), config()), + ).toBeNull(); }); it("includes the configured profile", () => { const cfg = config({ getAws: () => ({ region: "eu-west-1", profile: "work" }) }); - expect(shapeCredentials(AWS, awsToken, cfg)).toMatchObject({ + expect(shapeCredentials("bedrock", AWS, awsToken, cfg)).toMatchObject({ type: "aws-credentials", region: "eu-west-1", profile: "work", @@ -114,7 +148,7 @@ describe("shapeCredentials — AWS credentials JSON", () => { describe("shapeCredentials — Google Cloud credentials JSON", () => { it("parses project/location/token for a brokered token", () => { const token = JSON.stringify({ project: "p", location: "us-central1", token: "gcp-tok" }); - expect(shapeCredentials(GOOGLE, token, config())).toEqual({ + expect(shapeCredentials("google-vertex", GOOGLE, token, config())).toEqual({ type: "google-cloud", project: "p", location: "us-central1", @@ -124,7 +158,7 @@ describe("shapeCredentials — Google Cloud credentials JSON", () => { it("omits accessToken for the ADC fallback when no token is present", () => { const token = JSON.stringify({ project: "p", location: "us-central1" }); - expect(shapeCredentials(GOOGLE, token, config())).toEqual({ + expect(shapeCredentials("google-vertex", GOOGLE, token, config())).toEqual({ type: "google-cloud", project: "p", location: "us-central1", @@ -132,22 +166,27 @@ describe("shapeCredentials — Google Cloud credentials JSON", () => { }); it("returns null for a non-JSON token", () => { - expect(shapeCredentials(GOOGLE, "not-json", config())).toBeNull(); + expect(shapeCredentials("google-vertex", GOOGLE, "not-json", config())).toBeNull(); }); it("returns null when project or location is missing", () => { - expect(shapeCredentials(GOOGLE, JSON.stringify({ project: "p" }), config())).toBeNull(); - expect(shapeCredentials(GOOGLE, JSON.stringify({ location: "l" }), config())).toBeNull(); + expect( + shapeCredentials("google-vertex", GOOGLE, JSON.stringify({ project: "p" }), config()), + ).toBeNull(); + expect( + shapeCredentials("google-vertex", GOOGLE, JSON.stringify({ location: "l" }), config()), + ).toBeNull(); }); }); describe("shapeCredentials — apikey baseUrl + customHeaders", () => { it("reads baseUrl and customHeaders under the provider configKey", () => { const cfg = config({ - getBaseUrl: (k) => (k === "anthropic" ? "https://proxy" : undefined), - getCustomHeaders: (k) => (k === "anthropic" ? { "x-tenancy": "t" } : undefined), + getBaseUrl: ({ configKey }) => (configKey === "anthropic" ? "https://proxy" : undefined), + getCustomHeaders: ({ configKey }) => + configKey === "anthropic" ? { "x-tenancy": "t" } : undefined, }); - expect(shapeCredentials(ANTHROPIC, "sk", cfg)).toEqual({ + expect(shapeCredentials("anthropic", ANTHROPIC, "sk", cfg)).toEqual({ type: "apikey", apiKey: "sk", baseUrl: "https://proxy", @@ -157,16 +196,104 @@ describe("shapeCredentials — apikey baseUrl + customHeaders", () => { it("normalizes an empty customHeaders object to undefined", () => { expect( - shapeCredentials(ANTHROPIC, "sk", config({ getCustomHeaders: () => ({}) })), + shapeCredentials("anthropic", ANTHROPIC, "sk", config({ getCustomHeaders: () => ({}) })), ).toMatchObject({ customHeaders: undefined }); }); it("uses the authProviderId as configKey when no override exists (openai-api)", () => { const cfg = config({ - getCustomHeaders: (k) => (k === "openai-api" ? { "x-flag": "1" } : undefined), + getCustomHeaders: ({ configKey }) => + configKey === "openai-api" ? { "x-flag": "1" } : undefined, }); - expect(shapeCredentials(OPENAI, "sk", cfg)).toMatchObject({ + expect(shapeCredentials("openai", OPENAI, "sk", cfg)).toMatchObject({ customHeaders: { "x-flag": "1" }, }); }); }); + +// A `providers.custom` entry's id is the user's chosen name, so shaping can't +// recognize it by id: the readers have to be told *which* provider is being +// resolved, and structured base-URL derivation has to be declared on the +// mapping. Without both, a named `type: "aws"` / `type: "snowflake"` entry +// inherits bedrock's region or loses its Cortex URL. +describe("shapeCredentials — providers.custom entries", () => { + const CUSTOM_AWS = { + authProviderId: "my-bedrock", + credentialType: "aws-credentials", + } as const; + const CUSTOM_SNOWFLAKE = { + authProviderId: "my-snow", + credentialType: "apikey", + structuredBaseUrl: "snowflake", + } as const; + const CUSTOM_GATEWAY = { authProviderId: "my-gateway", credentialType: "apikey" } as const; + + const awsToken = JSON.stringify({ accessKeyId: "AKIA", secretAccessKey: "secret" }); + + it("asks getAws for the custom entry's own key, not the built-in bedrock one", () => { + const cfg = config({ + getAws: ({ providerId }) => + providerId === "my-bedrock" ? { region: "ca-central-1" } : { region: "us-west-2" }, + }); + expect(shapeCredentials("my-bedrock", CUSTOM_AWS, awsToken, cfg)).toMatchObject({ + region: "ca-central-1", + }); + }); + + it("derives the Cortex URL from the custom entry's own host", () => { + const cfg = config({ + getSnowflake: ({ providerId }) => + providerId === "my-snow" ? { host: "mine.snowflakecomputing.com" } : undefined, + }); + expect(shapeCredentials("my-snow", CUSTOM_SNOWFLAKE, "tok", cfg)).toMatchObject({ + baseUrl: "https://mine.snowflakecomputing.com/api/v2/cortex/v1", + }); + }); + + it("resolves a custom entry's flat baseUrl when it has no structured fields", () => { + const cfg = config({ + getBaseUrl: ({ providerId }) => + providerId === "my-snow" ? "https://mine.example.com/api/v2/cortex/v1" : undefined, + }); + expect(shapeCredentials("my-snow", CUSTOM_SNOWFLAKE, "tok", cfg)).toMatchObject({ + baseUrl: "https://mine.example.com/api/v2/cortex/v1", + }); + }); + + it("leaves a plain custom entry on the baseUrl path", () => { + // The declaration is what selects structured derivation: an entry that + // doesn't declare it must not pick up another provider's Snowflake config. + const cfg = config({ + getBaseUrl: () => "https://gateway.example.com", + getSnowflake: () => ({ host: "someone-else.snowflakecomputing.com" }), + }); + expect(shapeCredentials("my-gateway", CUSTOM_GATEWAY, "sk", cfg)).toMatchObject({ + baseUrl: "https://gateway.example.com", + }); + }); + + // A custom entry may legally be named `snowflake`, which is also the configKey + // `snowflake-cortex` derives. Two providers, one configKey: readers can only + // tell them apart by provider id. + it("distinguishes an entry named `snowflake` from the built-in snowflake-cortex", () => { + const NAMED_SNOWFLAKE = { + authProviderId: "snowflake", + credentialType: "apikey", + structuredBaseUrl: "snowflake", + } as const; + const hosts: Record = { + snowflake: "mine.snowflakecomputing.com", + "snowflake-cortex": "corp.snowflakecomputing.com", + }; + const cfg = config({ + getSnowflake: ({ providerId }) => ({ host: hosts[providerId] }), + }); + + expect(shapeCredentials("snowflake", NAMED_SNOWFLAKE, "tok", cfg)).toMatchObject({ + baseUrl: "https://mine.snowflakecomputing.com/api/v2/cortex/v1", + }); + expect(shapeCredentials("snowflake-cortex", SNOWFLAKE, "tok", cfg)).toMatchObject({ + baseUrl: "https://corp.snowflakecomputing.com/api/v2/cortex/v1", + }); + }); +}); diff --git a/packages/ai-credentials/src/types/credential-shaping.ts b/packages/ai-credentials/src/types/credential-shaping.ts index d726d12..fb55f0d 100644 --- a/packages/ai-credentials/src/types/credential-shaping.ts +++ b/packages/ai-credentials/src/types/credential-shaping.ts @@ -26,17 +26,42 @@ import { normalizeDatabricksHost, } from "./utils.js"; +/** + * Which structured connection fields a provider builds its base URL from, for + * the providers that don't use a plain `baseUrl`. + */ +export type StructuredBaseUrlSource = "snowflake" | "databricks"; + /** * Maps a provider to its auth extension registration and credential type. - * Subset of the full mapping — shaping only needs these two fields. + * Subset of the full mapping — shaping only needs these fields. */ export interface AuthProviderMapping { authProviderId: string; scopes: string[]; fallbackScopes?: string[][]; credentialType: "apikey" | "oauth" | "aws-credentials" | "google-cloud"; + /** + * Set for `providers.custom` entries whose base URL comes from structured + * fields. Built-in providers are recognized by their auth provider id + * ({@link BUILTIN_STRUCTURED_BASE_URL}); a custom entry's id is the user's + * chosen name, so its mapping has to say so. + */ + structuredBaseUrl?: StructuredBaseUrlSource; } +/** + * Built-in providers whose `apikey` base URL is derived from structured + * connection fields rather than a plain `baseUrl`. + * + * Kept here rather than in the injected mapping so hosts that hand-build + * mappings for built-ins keep their derivation. + */ +const BUILTIN_STRUCTURED_BASE_URL: Record = { + "snowflake-cortex": "snowflake", + databricks: "databricks", +}; + /** * Auth provider ID -> VS Code settings config section. * Most providers use the auth provider ID directly; legacy `anthropic-api` maps to `anthropic`. @@ -47,39 +72,75 @@ export const CONFIG_KEY_OVERRIDES: Record = { "snowflake-cortex": "snowflake", }; +/** + * Which provider a {@link CredentialConfig} read is for. + * + * Two fields because the two jobs need different keys, and only one of them is + * unique: + * + * - `providerId` is the identity. A built-in provider id or a `providers.custom` + * entry id, and the two spaces can't collide because custom names reserve + * every built-in id. Catalog-backed adapters answer from this. + * - `configKey` is the settings namespace (`authentication..*`) that + * settings-backed adapters read. It is **not** unique: `snowflake-cortex` + * derives the configKey `snowflake`, which is itself a legal custom entry + * name, so a reader that identifies a provider by configKey alone will hand a + * custom entry the built-in's connection. + * + * Passed as an object rather than two string parameters so the two can't be + * transposed silently. + */ +export interface CredentialConfigTarget { + readonly providerId: string; + readonly configKey: string; +} + /** * Reads the provider-extra config that shaping needs, abstracted over the * config source. Hosts inject catalog-backed adapters (reading the resolved * provider catalog's connection fields); Positron's renderer adapter reads its * own `IConfigurationService`. The shaper owns *which* keys to read (via - * `configKey`) so neither caller has to. + * {@link CredentialConfigTarget}) so neither caller has to. + * + * Every reader must answer for the requested provider only. A `providers.custom` + * entry has its own connection, so a host serving custom entries cannot answer + * from a fixed built-in provider: a named `type: "aws"` entry would inherit + * `bedrock`'s region. */ export interface CredentialConfig { /** `authentication..baseUrl` (the shaper normalizes empty -> undefined). */ - getBaseUrl(configKey: string): string | undefined; + getBaseUrl(target: CredentialConfigTarget): string | undefined; /** `authentication..customHeaders`. */ - getCustomHeaders(configKey: string): Record | undefined; + getCustomHeaders(target: CredentialConfigTarget): Record | undefined; /** AWS region/profile, from the resolved catalog's `connection.aws`. */ - getAws(): { region?: string; profile?: string } | undefined; + getAws(target: CredentialConfigTarget): { region?: string; profile?: string } | undefined; /** Snowflake host/account (`authentication.snowflake.credentials`, env on the bridge side). */ - getSnowflake(): { host?: string; account?: string } | undefined; + getSnowflake(target: CredentialConfigTarget): { host?: string; account?: string } | undefined; /** Databricks workspace host (`authentication.databricks.credentials`, env on the bridge side). */ - getDatabricks(): { host?: string } | undefined; + getDatabricks(target: CredentialConfigTarget): { host?: string } | undefined; } /** * Shape an already-resolved auth token into {@link ProviderCredentials}, or * `null` when the token cannot yield usable credentials (malformed JSON, missing * required fields). The mapping supplies the credential type and the auth - * provider id (from which the settings `configKey` is derived); `config` reads - * the provider-extra settings. + * provider id (from which the settings `configKey` is derived); `providerId` + * identifies which provider is being resolved, since the derived configKey does + * not (see {@link CredentialConfigTarget}); `config` reads the provider-extra + * settings. */ export function shapeCredentials( - mapping: Pick, + providerId: string, + mapping: Pick, rawToken: string, config: CredentialConfig, logger?: Logger, ): ProviderCredentials | null { + const target: CredentialConfigTarget = { + providerId, + configKey: CONFIG_KEY_OVERRIDES[mapping.authProviderId] ?? mapping.authProviderId, + }; + switch (mapping.credentialType) { case "oauth": return { type: "oauth", accessToken: rawToken }; @@ -122,7 +183,7 @@ export function shapeCredentials( return null; } // Region is not in the session -- the adapter resolves settings/env, default us-east-1. - const aws = config.getAws(); + const aws = config.getAws(target); return { type: "aws-credentials", region: aws?.region || "us-east-1", @@ -134,31 +195,45 @@ export function shapeCredentials( } case "apikey": { - const configKey = CONFIG_KEY_OVERRIDES[mapping.authProviderId] ?? mapping.authProviderId; - let baseUrl: string | undefined; - if (mapping.authProviderId === "snowflake-cortex") { - // Snowflake URL is built from host (preferred, for private-link/RCR) or account name. - const snowflake = config.getSnowflake(); - if (snowflake?.host) { - baseUrl = buildSnowflakeCortexUrlFromHost(snowflake.host); - } else if (snowflake?.account) { - baseUrl = buildSnowflakeCortexUrl(snowflake.account); + switch (mapping.structuredBaseUrl ?? BUILTIN_STRUCTURED_BASE_URL[mapping.authProviderId]) { + case "snowflake": { + // A flat `baseUrl` wins over the structured fields — the same + // precedence as the Node catalog paths + // (`conn.baseUrl ?? deriveSnowflakeBaseUrl(conn)`), so one + // providers.json can't route different hosts to different + // endpoints. The flat form is what standalone's + // Add-custom-provider form writes in custom-URL mode, and the + // only shape that can express a non-standard Cortex path. + // Otherwise the URL is built from host (preferred, for + // private-link/RCR) or account name. + const flat = config.getBaseUrl(target) || undefined; + const snowflake = flat ? undefined : config.getSnowflake(target); + if (flat) { + baseUrl = flat; + } else if (snowflake?.host) { + baseUrl = buildSnowflakeCortexUrlFromHost(snowflake.host); + } else if (snowflake?.account) { + baseUrl = buildSnowflakeCortexUrl(snowflake.account); + } + break; } - } 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); + case "databricks": { + // Databricks workspace host, with env fallback for managed environments + // (e.g. Posit Workbench injecting DATABRICKS_HOST into sessions). + const databricks = config.getDatabricks(target); + if (databricks?.host) { + baseUrl = normalizeDatabricksHost(databricks.host); + } + break; } - } else { - baseUrl = config.getBaseUrl(configKey) || undefined; + default: + baseUrl = config.getBaseUrl(target) || undefined; } // customHeaders share the `authentication.` namespace with // baseUrl. Empty objects normalize to undefined to match the pipeline. - const customHeadersRaw = config.getCustomHeaders(configKey); + const customHeadersRaw = config.getCustomHeaders(target); const customHeaders = customHeadersRaw && Object.keys(customHeadersRaw).length > 0 ? customHeadersRaw : undefined; diff --git a/packages/ai-credentials/src/types/index.ts b/packages/ai-credentials/src/types/index.ts index baa1aa4..1af6c3d 100644 --- a/packages/ai-credentials/src/types/index.ts +++ b/packages/ai-credentials/src/types/index.ts @@ -23,7 +23,12 @@ export type { } from "./credentials.js"; export { CONFIG_KEY_OVERRIDES, shapeCredentials } from "./credential-shaping.js"; -export type { AuthProviderMapping, CredentialConfig } from "./credential-shaping.js"; +export type { + AuthProviderMapping, + CredentialConfig, + CredentialConfigTarget, + StructuredBaseUrlSource, +} from "./credential-shaping.js"; export type { Logger } from "./logger.js";