Positron: support custom providers

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<string>();
 
+	// 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<string, number>();
+	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<ProviderCredentials | null> {
-		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<string[]>();
 
-	// Reverse map: auth provider id -> logical provider ids.
-	const authToLogical = new Map<string, string[]>();
-	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 {

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<string, AuthProviderMapping> = {
 	anthropic: { authProviderId: "anthropic-api", scopes: [], credentialType: "apikey" },
@@ -82,10 +83,13 @@ function testConfig(overrides: Partial<CredentialConfig> = {}): CredentialConfig
 	};
 }
 
-function makeBackend(configOverrides: Partial<CredentialConfig> = {}) {
+function makeBackend(
+	configOverrides: Partial<CredentialConfig> = {},
+	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"]]);
+		});
+	});
 });

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();