@adeildo/pi-providers 5.2.0 → 5.4.0

package/src/accounts/lift.ts

@@ -1,6 +1,7 @@
-// A provider with accounts is re-registered with its stream wrapped: each request resolves the
-// credential of the active account and injects it, so one provider keeps one model list. The auth
-// methods of the native provider do the work, so this file never speaks OAuth.
+// A provider with accounts is re-registered so the accounts own it: every way pi asks for auth answers
+// with the account in use, and each stream injects it again so a failure can move to the next one. One
+// provider keeps one model list. The auth methods of the native provider do the work, so this file
+// never speaks OAuth.
 import {
 	defaultProviderAuthContext,
 	lazyStream,
@@ -16,6 +17,7 @@ import {
 	type OAuthCredential,
 	type Provider,
 	type ProviderAuth,
+	type ProviderAuthInteraction,
 	type ProviderEnv,
 	type ProviderHeaders,
 	type SimpleStreamOptions,
@@ -23,15 +25,41 @@ import {
 	type TranscriptContext,
 } from "@earendil-works/pi-ai";
 
+import { classifyFailure, type Failure } from "../errors.ts";
+import { needsRefresh, refreshAccount, type Renewal } from "./refresh.ts";
+
+/** An account as a request sees it. */
+export interface AccountCredential {
+	id: string;
+	credential: Credential;
+}
+
 /** The account the next request uses, and where a refreshed token is kept. */
 export interface AccountSession {
-	resolve(): { id: string; credential: Credential } | undefined;
+	/** The account in use, or undefined once the provider has none left. */
+	resolve(): AccountCredential | undefined;
+	/** A subscription account, for when pi's copy needs a refresh and the one in use is a key. */
+	subscription?(): AccountCredential | undefined;
 	save(id: string, credential: Credential): void;
 	/** The account to try after this one hit a limit, or undefined to surface the error. */
-	afterLimit(
-		currentId: string,
-		detail: string,
-	): Promise<{ id: string; credential: Credential } | undefined>;
+	afterLimit(currentId: string, detail: string): Promise<AccountCredential | undefined>;
+	/** The credential was refused, so sign in again or try another account. */
+	afterAuthFailure?(currentId: string, detail: string): Promise<AccountCredential | undefined>;
+	/** Pi's own login just produced a credential; the feature keeps it under a name. */
+	adopt?(credential: Credential): Promise<void>;
+	/** A framed screen for a login, when the caller can draw one. */
+	loginScreen?(
+		label: string,
+		signal: AbortSignal,
+	): Promise<{ interaction: ProviderAuthInteraction; close(): void } | undefined>;
+	/** The credential on disk now, in case another pi refreshed it since we read ours. */
+	freshen?(id: string): Credential | undefined;
+	/** The file the refresh takes across processes, when the caller shares one. */
+	lockPath?(): string;
+	/** The quota headers of a response, kept for the account that served it. */
+	noteUsage?(id: string, headers: Record<string, string>): void;
+	/** The account served a request, so a limit it had is over. */
+	served?(id: string): void;
 }
 
 /** The native provider behind a lifted one, so lifting twice wraps the same object. */
@@ -60,7 +88,7 @@ export function liftProvider(provider: Provider, session: AccountSession): Provi
 		name: provider.name,
 		...(provider.baseUrl === undefined ? {} : { baseUrl: provider.baseUrl }),
 		...(provider.headers === undefined ? {} : { headers: provider.headers }),
-		auth: liftedAuth(provider, session),
+		auth: liftedAuth(provider, session, refreshing),
 		getModels: () => provider.getModels(),
 		...(getAllModels === undefined ? {} : { getAllModels: () => getAllModels() }),
 		...(refreshModels === undefined ? {} : { refreshModels: (context) => refreshModels(context) }),
@@ -109,20 +137,15 @@ function accountStream(
 	options?: Options,
 ): AssistantMessageEventStream {
 	const first = session.resolve();
-	if (first === undefined) return call(provider, kind, model, context, options ?? {});
-
 	const settled = options ?? {};
 	return lazyStream(model, async () =>
 		attempts(provider, session, refreshing, kind, model, context, settled, first),
 	);
 }
 
-/** One account may serve a limit that another one does not have. */
+/** One account may serve a limit that another one does not have. A re-login costs an attempt too. */
 const MAX_ATTEMPTS = 3;
 
-// A quota or a rate limit belongs to the account, so another one can serve. A busy server cannot.
-const LIMIT_PATTERN = /quota|rate[ _-]?limit|too many requests|usage limit|429/i;
-
 async function* attempts(
 	provider: Provider,
 	session: AccountSession,
@@ -131,16 +154,11 @@ async function* attempts(
 	model: Model<Api>,
 	context: TranscriptContext,
 	options: Options,
-	first: { id: string; credential: Credential },
+	first: AccountCredential | undefined,
 ): AsyncGenerator<AssistantMessageEvent> {
 	let account = first;
 
 	for (let attempt = 1; ; attempt++) {
-		// oxlint-disable-next-line no-await-in-loop -- an attempt starts only after the one before it.
-		const resolved = await resolveAuth(provider, session, refreshing, account, options.signal);
-		const withAuth = applyAuth(resolved, model, options);
-		const source = call(provider, kind, withAuth.model, context, withAuth.options);
-
 		// Hold the events until something is visible: a retry on another account must not repeat a
 		// message the caller already saw.
 		const buffered: AssistantMessageEvent[] = [];
@@ -148,6 +166,15 @@ async function* attempts(
 		let failure: AssistantMessage | undefined;
 		let thrown: unknown;
 		try {
+			// The auth resolution is inside the try too: a refresh that fails is an auth failure, and it
+			// must be able to move to the next account like any other.
+			// oxlint-disable-next-line no-await-in-loop -- an attempt starts only after the one before it.
+			const withAuth = await authFor(provider, session, refreshing, account, model, options);
+			const source = call(provider, kind, withAuth.model, context, {
+				...withAuth.options,
+				onResponse: noteUsage(withAuth.options.onResponse, session, account?.id),
+			});
+
 			// oxlint-disable-next-line no-await-in-loop -- a stream is read one event at a time.
 			for await (const event of source) {
 				if (isContent(event)) {
@@ -167,15 +194,15 @@ async function* attempts(
 			thrown = error;
 		}
 
-		if (!output && attempt < MAX_ATTEMPTS && failure !== undefined) {
-			const text = failure.errorMessage;
-			if (text !== undefined && LIMIT_PATTERN.test(text)) {
-				// oxlint-disable-next-line no-await-in-loop -- the answer comes from the user.
-				const next = await session.afterLimit(account.id, text);
-				if (next !== undefined) {
-					account = next;
-					continue;
-				}
+		if (failure === undefined && thrown === undefined && account !== undefined)
+			session.served?.(account.id);
+
+		if (!output && attempt < MAX_ATTEMPTS) {
+			// oxlint-disable-next-line no-await-in-loop -- the answer comes from the user.
+			const next = await moveOn(session, account?.id, failureOf(failure, thrown, model.provider));
+			if (next !== undefined) {
+				account = next;
+				continue;
 			}
 		}
 
@@ -185,10 +212,59 @@ async function* attempts(
 	}
 }
 
+/** The classified failure behind an attempt: the provider's own error, or what it threw. */
+function failureOf(
+	failed: AssistantMessage | undefined,
+	thrown: unknown,
+	provider: string,
+): Failure | undefined {
+	if (failed?.errorMessage !== undefined) return classifyFailure(failed.errorMessage, provider);
+	return thrown === undefined ? undefined : classifyFailure(thrown, provider);
+}
+
+/** What the policy says after a failure: another account, a sign-in, or nothing. */
+async function moveOn(
+	session: AccountSession,
+	accountId: string | undefined,
+	failure: Failure | undefined,
+): Promise<AccountCredential | undefined> {
+	// Without an account there is nothing to move from, and the error is the provider's to show.
+	if (accountId === undefined) return undefined;
+	if (failure?.kind === "usage") return await session.afterLimit(accountId, failure.message);
+	if (failure?.kind === "auth") return await session.afterAuthFailure?.(accountId, failure.message);
+	return undefined;
+}
+
 function isContent(event: AssistantMessageEvent): boolean {
 	return event.type !== "start" && event.type !== "done" && event.type !== "error";
 }
 
+/** Pi hands the response headers here, so the account that served it keeps its quota current. */
+function noteUsage(
+	next: Options["onResponse"],
+	session: AccountSession,
+	accountId: string | undefined,
+): Options["onResponse"] {
+	return async (response, model) => {
+		if (accountId !== undefined) session.noteUsage?.(accountId, response.headers);
+		await next?.(response, model);
+	};
+}
+
+/** The request auth for an attempt: the account's, or untouched once the provider has none. */
+async function authFor(
+	provider: Provider,
+	session: AccountSession,
+	refreshing: Map<string, Promise<OAuthCredential>>,
+	account: AccountCredential | undefined,
+	model: Model<Api>,
+	options: Options,
+): Promise<{ model: Model<Api>; options: Options }> {
+	if (account === undefined) return { model, options };
+	const resolved = await resolveAuth(provider, session, refreshing, account, options.signal);
+	return applyAuth(resolved, model, options);
+}
+
 function call(
 	provider: Provider,
 	kind: "stream" | "streamSimple",
@@ -210,7 +286,7 @@ async function resolveAuth(
 	provider: Provider,
 	session: AccountSession,
 	refreshing: Map<string, Promise<OAuthCredential>>,
-	account: { id: string; credential: Credential },
+	account: AccountCredential,
 	signal: AbortSignal | undefined,
 ): Promise<ResolvedAuth> {
 	const abort = signal ?? new AbortController().signal;
@@ -238,10 +314,7 @@ async function resolveAuth(
 	return resolved === undefined ? { auth: {} } : { auth: resolved.auth, env: resolved.env };
 }
 
-// Pi refreshes its own token under the store lock, five minutes early. Ours is ours to refresh, and
-// one refresh per account is shared, so two requests cannot rotate the token twice.
-const REFRESH_MARGIN_MS = 5 * 60 * 1000;
-
+// One refresh per account is shared, so two requests cannot rotate the token twice.
 async function refreshToken(
 	oauth: OAuthAuth,
 	session: AccountSession,
@@ -250,15 +323,16 @@ async function refreshToken(
 	credential: OAuthCredential,
 	signal: AbortSignal,
 ): Promise<OAuthCredential> {
-	if (credential.expires > Date.now() + REFRESH_MARGIN_MS) return credential;
+	if (!needsRefresh(credential)) return credential;
 
 	const inFlight = refreshing.get(accountId);
 	if (inFlight !== undefined) return await inFlight;
 
-	const task = oauth.refresh(credential, signal).then((fresh) => {
-		session.save(accountId, fresh);
-		return fresh;
-	});
+	const renewal: Renewal = {
+		freshen: (id) => session.freshen?.(id),
+		save: (id, fresh) => session.save(id, fresh),
+	};
+	const task = refreshAccount(oauth, renewal, accountId, credential, signal, session.lockPath?.());
 	refreshing.set(accountId, task);
 	try {
 		return await task;
@@ -306,26 +380,104 @@ function mergeHeaders(
 	return merged;
 }
 
-// The runtime reads auth from the provider, so it has to report configured when an account exists,
-// even though the credential itself is injected into the request by `accountStream`. Reporting the
-// account's own type keeps `isUsingOAuth` honest, which is what bills a subscription request.
-function liftedAuth(provider: Provider, session: AccountSession): ProviderAuth {
+// Pi resolves a request's auth from the credential it stored for the provider, before the stream runs.
+// Every one of those paths answers with the account in use here, so a classifier, an image or a
+// deferred call bills the same account a chat does, and pi never refreshes its own copy: when that
+// copy expires, the refresh renews an account and hands pi the result. Pi keeps that copy only because
+// it decides at startup, before any extension runs, which providers are logged in.
+function liftedAuth(
+	provider: Provider,
+	session: AccountSession,
+	refreshing: Map<string, Promise<OAuthCredential>>,
+): ProviderAuth {
 	const original = provider.auth.apiKey;
 	const login = original?.login;
-	return {
+	const oauth = provider.auth.oauth;
+	const auth: ProviderAuth = {
 		apiKey: {
 			name: original?.name ?? `${provider.name} accounts`,
-			...(login === undefined ? {} : { login: (interaction) => login(interaction) }),
+			...(login === undefined
+				? {}
+				: { login: (interaction) => adopted(login, interaction, session, provider.name) }),
+			// Reporting the account's own type keeps `isUsingOAuth` honest, which bills a subscription.
 			async check(input) {
 				const account = session.resolve();
 				if (account !== undefined) return { type: account.credential.type, source: "account" };
 				return original?.check?.(input);
 			},
 			async resolve(input) {
-				if (session.resolve() !== undefined) return { auth: {}, source: "account" };
-				return original?.resolve(input);
+				const resolved = await authInUse(provider, session, refreshing, input.signal);
+				if (resolved === undefined) return original?.resolve(input);
+				return { ...resolved, source: "account" };
 			},
 		},
-		...(provider.auth.oauth === undefined ? {} : { oauth: provider.auth.oauth }),
 	};
+	if (oauth !== undefined) {
+		auth.oauth = {
+			name: oauth.name,
+			...(oauth.isSubscription === undefined ? {} : { isSubscription: oauth.isSubscription }),
+			...(oauth.loginLabel === undefined ? {} : { loginLabel: oauth.loginLabel }),
+			login: (interaction, options) =>
+				adopted((next) => oauth.login(next, options), interaction, session, provider.name),
+			async refresh(credential, signal) {
+				const account = subscriptionOf(session);
+				if (account === undefined) return await oauth.refresh(credential, signal);
+				return await refreshToken(
+					oauth,
+					session,
+					refreshing,
+					account.id,
+					account.credential,
+					signal,
+				);
+			},
+			async toAuth(credential) {
+				const signal = new AbortController().signal;
+				const resolved = await authInUse(provider, session, refreshing, signal);
+				return resolved?.auth ?? (await oauth.toAuth(credential));
+			},
+		};
+	}
+	return auth;
+}
+
+/** The auth of the account in use, or undefined once the provider has none. */
+async function authInUse(
+	provider: Provider,
+	session: AccountSession,
+	refreshing: Map<string, Promise<OAuthCredential>>,
+	signal: AbortSignal,
+): Promise<ResolvedAuth | undefined> {
+	const account = session.resolve();
+	if (account === undefined) return undefined;
+	return await resolveAuth(provider, session, refreshing, account, signal);
+}
+
+/** The OAuth account pi's copy is renewed from: the one in use, else any subscription. */
+function subscriptionOf(
+	session: AccountSession,
+): { id: string; credential: OAuthCredential } | undefined {
+	const candidates = [session.resolve(), session.subscription?.()];
+	for (const account of candidates) {
+		if (account?.credential.type === "oauth")
+			return { id: account.id, credential: account.credential };
+	}
+	return undefined;
+}
+
+/** Runs a login on our framed screen when there is one, then lets the feature keep the credential. */
+async function adopted<T extends Credential>(
+	login: (interaction: ProviderAuthInteraction) => Promise<T>,
+	interaction: ProviderAuthInteraction,
+	session: AccountSession,
+	label: string,
+): Promise<T> {
+	const screen = await session.loginScreen?.(label, interaction.signal);
+	try {
+		const credential = await login(screen?.interaction ?? interaction);
+		await session.adopt?.(credential);
+		return credential;
+	} finally {
+		screen?.close();
+	}
 }

package/src/accounts/lock.ts

@@ -0,0 +1,54 @@
+// Serializes the refresh across pi processes. A lock older than its deadline belonged to a process
+// that died, and is taken over.
+import { statSync, unlinkSync, writeFileSync } from "node:fs";
+
+const STALE_MS = 60_000;
+const RETRY_MS = 50;
+const WAIT_MS = 5_000;
+
+/** Runs `work` holding the lock, or without it when it cannot be taken in time. */
+export async function withLock<T>(
+	path: string,
+	work: () => Promise<T>,
+	waitMs = WAIT_MS,
+): Promise<T> {
+	const deadline = Date.now() + waitMs;
+	while (!take(path)) {
+		if (Date.now() >= deadline) return await work();
+		// oxlint-disable-next-line no-await-in-loop -- the lock is retried until its deadline.
+		await new Promise((resolve) => setTimeout(resolve, RETRY_MS));
+	}
+	try {
+		return await work();
+	} finally {
+		release(path);
+	}
+}
+
+function take(path: string): boolean {
+	try {
+		// Opens, writes and closes in one call: handing writeSync's return value to closeSync closed
+		// the byte count, and for a 5 or 6 digit pid that count is a descriptor libuv owns.
+		writeFileSync(path, String(process.pid), { flag: "wx" });
+		return true;
+	} catch {
+		dropStale(path);
+		return false;
+	}
+}
+
+function dropStale(path: string): void {
+	try {
+		if (Date.now() - statSync(path).mtimeMs > STALE_MS) unlinkSync(path);
+	} catch {
+		// The holder released it between the open and the stat, which is the good case.
+	}
+}
+
+function release(path: string): void {
+	try {
+		unlinkSync(path);
+	} catch {
+		// Already gone, which is also the point.
+	}
+}

package/src/accounts/login.ts

@@ -4,6 +4,20 @@ import type {
 	Provider,
 	ProviderAuthInteraction,
 } from "@earendil-works/pi-ai";
+import { readStoredCredential } from "@earendil-works/pi-coding-agent";
+
+/**
+ * The credential pi's own `/login` wrote for a provider. Read once, when the provider gets its first
+ * account, so that login becomes an account too. After that the store is the only one consulted.
+ */
+export function piLogin(providerId: string): Credential | undefined {
+	return readStoredCredential(providerId);
+}
+
+/** The word a person reads for a credential type. */
+export function kindText(type: Credential["type"]): string {
+	return type === "oauth" ? "subscription" : "api key";
+}
 
 /** One way a provider can log in, named the way its own login dialog names it. */
 export interface LoginMethod {

package/src/accounts/pins.ts

@@ -5,54 +5,39 @@ import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
 
 import { SESSION_ENTRY } from "./names.ts";
 
-/** The account pinned for a provider. `null` means the credential pi itself has. */
-export type Pins = Map<string, string | null>;
+/** The account pinned for a provider. */
+export type Pins = Map<string, string>;
 
-export type PinEntry =
-	| { kind: "account"; provider: string; account: string }
-	| { kind: "default"; provider: string };
-
-export function record(pi: ExtensionAPI, entry: PinEntry): void {
-	pi.appendEntry(SESSION_ENTRY, entry);
+export interface PinEntry {
+	kind: "account";
+	provider: string;
+	account: string;
 }
 
 /** Pins the account for the rest of the session, and writes it into the session file. */
-export function pin(
-	pi: ExtensionAPI,
-	pins: Pins,
-	providerId: string,
-	accountId: string | null,
-): void {
+export function pin(pi: ExtensionAPI, pins: Pins, providerId: string, accountId: string): void {
 	pins.set(providerId, accountId);
-	record(
-		pi,
-		accountId === null
-			? { kind: "default", provider: providerId }
-			: { kind: "account", provider: providerId, account: accountId },
-	);
+	pi.appendEntry(SESSION_ENTRY, {
+		kind: "account",
+		provider: providerId,
+		account: accountId,
+	} satisfies PinEntry);
 }
 
 /** The pins a branch recorded, newest winning. */
 export function replay(branch: readonly unknown[]): Pins {
 	const pins: Pins = new Map();
 	for (const item of branch) {
-		const entry = pinEntry(item);
-		if (entry === undefined) continue;
-		if (entry.kind === "default") pins.set(entry.provider, null);
-		else pins.set(entry.provider, entry.account);
+		if (!isObject(item) || item.type !== "custom" || item.customType !== SESSION_ENTRY) continue;
+		const data = item.data;
+		if (!isObject(data) || typeof data.provider !== "string") continue;
+		if (data.kind === "account" && typeof data.account === "string") {
+			pins.set(data.provider, data.account);
+		} else {
+			// Older sessions pinned pi's own credential, which is no account any more: the session goes
+			// back to the store's choice.
+			pins.delete(data.provider);
+		}
 	}
 	return pins;
 }
-
-function pinEntry(item: unknown): PinEntry | undefined {
-	if (!isObject(item) || item.type !== "custom" || item.customType !== SESSION_ENTRY)
-		return undefined;
-	const data = item.data;
-	if (!isObject(data)) return undefined;
-	const provider = typeof data.provider === "string" ? data.provider : undefined;
-	if (provider === undefined) return undefined;
-	if (data.kind === "default") return { kind: "default", provider };
-	if (data.kind === "account" && typeof data.account === "string")
-		return { kind: "account", provider, account: data.account };
-	return undefined;
-}