@narumitw/pi-usage 0.52.0 → 0.52.2

package/src/codex-resets.ts

@@ -1,6 +1,10 @@
 import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
 import { readStoredCredential } from "@earendil-works/pi-coding-agent";
 import { fingerprintResolvedAuth, sanitizeDisplayText } from "./core.js";
+import {
+	fallbackOAuthCredentialCandidates,
+	type OAuthCredentialCandidateReader,
+} from "./oauth-credential-source.js";
 import {
 	AUTH_FINGERPRINT_SALT,
 	adapterForProvider,
@@ -98,6 +102,7 @@ export async function resolveCodexResetAuth(
 	ctx: ExtensionContext,
 	salt: Uint8Array = AUTH_FINGERPRINT_SALT,
 	credentialReader: StoredCredentialReader = readStoredCredential,
+	candidateReader?: OAuthCredentialCandidateReader,
 ): Promise<ResolvedUsageAuth> {
 	const model = ctx.model;
 	if (model?.provider !== "openai-codex") {
@@ -112,26 +117,22 @@ export async function resolveCodexResetAuth(
 	}
 	if (!auth) throw new Error("No runtime credential is configured for OpenAI Codex.");
 
-	const credential = asObject(credentialReader("openai-codex"));
-	if (credential?.type !== "oauth") {
-		throw new Error(
-			"Usage limit resets require the OpenAI Codex OAuth account configured through Pi /login.",
-		);
-	}
-	const storedAccess = asNonemptyString(credential.access);
 	const resolvedAccess = bearerToken(headerValue(auth.headers, "Authorization")) ?? auth.apiKey;
-	if (!storedAccess || !resolvedAccess) {
-		throw new Error("OpenAI Codex OAuth credentials were incomplete.");
-	}
-	if (storedAccess !== resolvedAccess) {
-		throw new Error(
-			"The active OpenAI Codex runtime account does not match Pi's stored OAuth account.",
-		);
+	if (!resolvedAccess) throw new Error("OpenAI Codex OAuth credentials were incomplete.");
+	const resolvedAccountId = codexAccountIdFromAccessToken(resolvedAccess);
+	if (!resolvedAccountId) {
+		throw new Error("The active OpenAI Codex access token did not contain a valid account ID.");
 	}
-	const accountId = validHeaderValue(credential.accountId);
-	if (!accountId)
-		throw new Error("The OpenAI Codex OAuth credential did not include a valid account ID.");
-
+	const offered = candidateReader
+		? candidateReader(ctx, "openai-codex")
+		: fallbackOAuthCredentialCandidates("openai-codex", credentialReader);
+	if (!offered.ok) throw new Error("OpenAI Codex OAuth credential discovery failed closed.");
+	const { accountId, storedAccess } = selectCodexResetCredential(
+		offered.candidates,
+		resolvedAccess,
+		resolvedAccountId,
+		offered.offeredCount === 0,
+	);
 	const authorization = `Bearer ${resolvedAccess}`;
 	const headers = {
 		Authorization: authorization,
@@ -226,6 +227,70 @@ export function normalizeCodexResetCreditsPayload(
 	return { availableCount, options };
 }
 
+function selectCodexResetCredential(
+	candidates: readonly unknown[],
+	resolvedAccess: string,
+	resolvedAccountId: string,
+	standaloneFallback: boolean,
+): { accountId: string; storedAccess: string } {
+	let sawOAuth = false;
+	let sawMatchingAccess = false;
+	let sawInvalidAccountId = false;
+	const matches = new Map<string, { accountId: string; storedAccess: string }>();
+	for (const candidate of candidates) {
+		try {
+			const credential = asObject(candidate);
+			if (credential?.type !== "oauth") continue;
+			sawOAuth = true;
+			const storedAccess = asNonemptyString(credential.access);
+			if (storedAccess !== resolvedAccess) continue;
+			sawMatchingAccess = true;
+			const accountId = validHeaderValue(credential.accountId);
+			const refresh = asNonemptyString(credential.refresh);
+			if (!accountId || accountId !== resolvedAccountId || !refresh) {
+				sawInvalidAccountId = true;
+				continue;
+			}
+			matches.set(refresh, { accountId, storedAccess });
+		} catch {
+			// Malformed candidates never authorize a reset request.
+		}
+	}
+	if (sawInvalidAccountId) {
+		throw new Error("The OpenAI Codex OAuth credential did not include a valid account ID.");
+	}
+	if (matches.size > 1) {
+		throw new Error("Conflicting OAuth credentials match the active OpenAI Codex runtime account.");
+	}
+	const match = matches.values().next().value;
+	if (match) return match;
+	if (!sawOAuth) {
+		throw new Error(
+			standaloneFallback
+				? "Usage limit resets require the OpenAI Codex OAuth account configured through Pi /login."
+				: "Usage limit resets require an OpenAI Codex OAuth account configured through Pi /login or a compatible credential source.",
+		);
+	}
+	if (sawMatchingAccess) throw new Error("OpenAI Codex OAuth credentials were incomplete.");
+	throw new Error(
+		standaloneFallback
+			? "The active OpenAI Codex runtime account does not match Pi's stored OAuth account."
+			: "The active OpenAI Codex runtime account does not match any available OAuth account.",
+	);
+}
+
+function codexAccountIdFromAccessToken(access: string): string | undefined {
+	try {
+		const parts = access.split(".");
+		if (parts.length !== 3 || !parts[1]) return undefined;
+		const payload = JSON.parse(Buffer.from(parts[1], "base64url").toString("utf8")) as unknown;
+		const claims = asObject(asObject(payload)?.["https://api.openai.com/auth"]);
+		return validHeaderValue(claims?.chatgpt_account_id);
+	} catch {
+		return undefined;
+	}
+}
+
 function normalizeResetOption(credit: Record<string, unknown>): CodexResetOption {
 	const creditId = asOpaqueId(credit.id);
 	if (!creditId) throw new Error("Codex reset credits response returned an invalid credit ID.");

package/src/oauth-credential-source.ts

@@ -0,0 +1,88 @@
+import type { OAuthCredential } from "@earendil-works/pi-ai";
+import {
+	type ExtensionAPI,
+	type ExtensionContext,
+	readStoredCredential,
+} from "@earendil-works/pi-coding-agent";
+
+export const OAUTH_CREDENTIAL_SOURCE_CHANNEL = "oauth:credential-source:v1";
+
+export type StoredCredentialReader = (providerId: string) => unknown;
+
+export type OAuthCredentialCandidates =
+	| { ok: true; candidates: readonly OAuthCredential[]; offeredCount?: number }
+	| { ok: false };
+
+export type OAuthCredentialCandidateReader = (
+	ctx: ExtensionContext,
+	providerId: string,
+) => OAuthCredentialCandidates;
+
+export function createOAuthCredentialCandidateReader(
+	pi: ExtensionAPI,
+	credentialReader: StoredCredentialReader = readStoredCredential,
+): OAuthCredentialCandidateReader {
+	return (ctx, providerId) =>
+		collectOAuthCredentialCandidates(pi, ctx, providerId, credentialReader);
+}
+
+export function collectOAuthCredentialCandidates(
+	pi: Pick<ExtensionAPI, "events">,
+	ctx: ExtensionContext,
+	providerId: string,
+	credentialReader: StoredCredentialReader = readStoredCredential,
+): OAuthCredentialCandidates {
+	const candidates: OAuthCredential[] = [];
+	let collecting = true;
+	const request = Object.freeze({
+		session: ctx.sessionManager,
+		provider: providerId,
+		offer(candidate: unknown) {
+			if (!collecting) return;
+			const clone = cloneOAuthCredential(candidate);
+			if (clone) candidates.push(clone);
+		},
+	});
+	try {
+		pi.events.emit(OAUTH_CREDENTIAL_SOURCE_CHANNEL, request);
+	} catch {
+		return { ok: false };
+	} finally {
+		collecting = false;
+	}
+	const offeredCount = candidates.length;
+	try {
+		const fallback = cloneOAuthCredential(credentialReader(providerId));
+		if (fallback) candidates.push(fallback);
+	} catch {
+		// A malformed or unavailable standalone credential is equivalent to no fallback.
+	}
+	return { ok: true, candidates, offeredCount };
+}
+
+export function fallbackOAuthCredentialCandidates(
+	providerId: string,
+	credentialReader: StoredCredentialReader,
+): OAuthCredentialCandidates {
+	try {
+		const credential = cloneOAuthCredential(credentialReader(providerId));
+		return { ok: true, candidates: credential ? [credential] : [], offeredCount: 0 };
+	} catch {
+		return { ok: true, candidates: [], offeredCount: 0 };
+	}
+}
+
+function cloneOAuthCredential(value: unknown): OAuthCredential | undefined {
+	try {
+		if (!value || typeof value !== "object" || Array.isArray(value)) return undefined;
+		const clone = structuredClone(value) as OAuthCredential;
+		if (!clone || typeof clone !== "object" || Array.isArray(clone)) return undefined;
+		if (clone.type !== "oauth") return undefined;
+		if (typeof clone.access !== "string" || !clone.access) return undefined;
+		if (typeof clone.refresh !== "string" || !clone.refresh) return undefined;
+		if (typeof clone.expires !== "number" || !Number.isFinite(clone.expires)) return undefined;
+		return clone;
+	} catch {
+		return undefined;
+	}
+}

package/src/query.ts

@@ -1,6 +1,10 @@
 import { randomBytes } from "node:crypto";
 import { type ExtensionContext, readStoredCredential } from "@earendil-works/pi-coding-agent";
 import { errorMessage, fingerprintResolvedAuth, redactUsageError } from "./core.js";
+import {
+	fallbackOAuthCredentialCandidates,
+	type OAuthCredentialCandidateReader,
+} from "./oauth-credential-source.js";
 import { normalizeCodexBackendPayload } from "./providers/codex.js";
 import { normalizeGitHubCopilotUsagePayload } from "./providers/github-copilot.js";
 import { normalizeOpenCodeZenPayload } from "./providers/opencode-zen.js";
@@ -111,6 +115,7 @@ export async function resolveUsageAuth(
 	adapter: UsageProviderAdapter,
 	salt: Uint8Array = AUTH_FINGERPRINT_SALT,
 	credentialReader: StoredCredentialReader = readStoredCredential,
+	candidateReader?: OAuthCredentialCandidateReader,
 ): Promise<ResolvedUsageAuth | undefined> {
 	if (ctx.model?.provider === adapter.id && !hasOfficialOrigin(ctx.model, adapter.id)) {
 		throw new Error(
@@ -144,7 +149,19 @@ export async function resolveUsageAuth(
 	const auth = modelAuth ?? providerResult?.auth;
 	if (!auth) return undefined;
 	if (adapter.id === "github-copilot") {
-		return resolveGitHubCopilotUsageAuth(auth, model, salt, credentialReader);
+		const offered = candidateReader
+			? candidateReader(ctx, adapter.id)
+			: fallbackOAuthCredentialCandidates(adapter.id, credentialReader);
+		if (!offered.ok) {
+			throw new Error("GitHub Copilot OAuth credential discovery failed closed.");
+		}
+		return resolveGitHubCopilotUsageAuth(
+			auth,
+			model,
+			salt,
+			offered.candidates,
+			offered.offeredCount === 0,
+		);
 	}
 	const authorization = authorizationFrom(auth);
 	if (!authorization) return undefined;
@@ -334,33 +351,73 @@ function resolveGitHubCopilotUsageAuth(
 	auth: RequestAuth,
 	model: PiModel,
 	salt: Uint8Array,
-	credentialReader: StoredCredentialReader,
+	candidates: readonly unknown[],
+	standaloneFallback: boolean,
 ): ResolvedUsageAuth {
-	const credential = asObject(credentialReader("github-copilot"));
-	if (credential?.type !== "oauth") {
-		throw new Error(
-			"GitHub Copilot usage requires the OAuth account configured through Pi /login.",
-		);
+	const resolvedAccess = bearerToken(headerValue(auth.headers, "Authorization")) ?? auth.apiKey;
+	if (!resolvedAccess) throw new Error("GitHub Copilot OAuth credentials were incomplete.");
+	let sawOAuth = false;
+	let sawMatchingAccess = false;
+	let sawIncompleteMatch = false;
+	let sawEnterpriseMatch = false;
+	const matches = new Map<string, { refresh: string; storedAccess: string }>();
+	for (const candidate of candidates) {
+		try {
+			const credential = asObject(candidate);
+			if (credential?.type !== "oauth") continue;
+			sawOAuth = true;
+			const storedAccess =
+				typeof credential.access === "string" && credential.access ? credential.access : undefined;
+			if (storedAccess !== resolvedAccess) continue;
+			sawMatchingAccess = true;
+			const enterpriseUrl = credential.enterpriseUrl;
+			if (
+				typeof enterpriseUrl === "string" &&
+				enterpriseUrl &&
+				!isPublicGitHubDomain(enterpriseUrl)
+			) {
+				sawEnterpriseMatch = true;
+				continue;
+			}
+			const refresh =
+				typeof credential.refresh === "string" && credential.refresh
+					? credential.refresh
+					: undefined;
+			if (!refresh) {
+				sawIncompleteMatch = true;
+				continue;
+			}
+			matches.set(`${storedAccess.length}:${storedAccess}${refresh}`, { refresh, storedAccess });
+		} catch {
+			// Malformed candidates never authorize a provider request.
+		}
 	}
-	if (
-		typeof credential.enterpriseUrl === "string" &&
-		credential.enterpriseUrl &&
-		!isPublicGitHubDomain(credential.enterpriseUrl)
-	) {
+	if (sawEnterpriseMatch) {
 		throw new Error("GitHub Copilot usage does not yet support GitHub Enterprise accounts.");
 	}
-	const refresh = typeof credential.refresh === "string" ? credential.refresh : undefined;
-	const storedAccess = typeof credential.access === "string" ? credential.access : undefined;
-	const resolvedAccess = bearerToken(headerValue(auth.headers, "Authorization")) ?? auth.apiKey;
-	if (!refresh || !storedAccess || !resolvedAccess) {
-		throw new Error("GitHub Copilot OAuth credentials were incomplete.");
+	if (sawIncompleteMatch) throw new Error("GitHub Copilot OAuth credentials were incomplete.");
+	if (matches.size > 1) {
+		throw new Error(
+			"Conflicting OAuth credentials match the active GitHub Copilot runtime account.",
+		);
 	}
-	if (storedAccess !== resolvedAccess) {
+	const match = matches.values().next().value;
+	if (!match) {
+		if (!sawOAuth) {
+			throw new Error(
+				standaloneFallback
+					? "GitHub Copilot usage requires the OAuth account configured through Pi /login."
+					: "GitHub Copilot usage requires an OAuth account configured through Pi /login or a compatible credential source.",
+			);
+		}
+		if (sawMatchingAccess) throw new Error("GitHub Copilot OAuth credentials were incomplete.");
 		throw new Error(
-			"The active GitHub Copilot runtime account does not match Pi's stored OAuth account.",
+			standaloneFallback
+				? "The active GitHub Copilot runtime account does not match Pi's stored OAuth account."
+				: "The active GitHub Copilot runtime account does not match any available OAuth account.",
 		);
 	}
-
+	const { refresh, storedAccess } = match;
 	const authorization = `Bearer ${refresh}`;
 	const headers = {
 		Authorization: authorization,

package/src/usage.ts

@@ -22,6 +22,7 @@ import {
 } from "./codex-resets.js";
 import { awaitWithDeadline, errorMessage, runWithConcurrency, UsageCache } from "./core.js";
 import { formatProviderStates, formatUsageStatusline } from "./format.js";
+import { createOAuthCredentialCandidateReader } from "./oauth-credential-source.js";
 import {
 	adapterForProvider,
 	isStaleExtensionContextError,
@@ -80,6 +81,7 @@ export default function usageExtension(
 	dependencies: UsageExtensionDependencies = {},
 ) {
 	const credentialReader = dependencies.credentialReader;
+	const credentialCandidates = createOAuthCredentialCandidateReader(pi, credentialReader);
 	const createRedemptionId = dependencies.createRedemptionId ?? randomUUID;
 	const settingsRuntime = dependencies.settingsRuntime ?? createUsageSettingsRuntime();
 	const cache = new UsageCache(CACHE_TTL_MS);
@@ -189,7 +191,7 @@ export default function usageExtension(
 		let auth: ResolvedUsageAuth | undefined;
 		try {
 			auth = await awaitWithDeadline(
-				resolveUsageAuth(ctx, adapter),
+				resolveUsageAuth(ctx, adapter, undefined, credentialReader, credentialCandidates),
 				signal,
 				DEFAULT_TIMEOUT_MS,
 				`resolving ${adapter.displayName} runtime auth`,
@@ -419,7 +421,7 @@ export default function usageExtension(
 			if (!adapter) return false;
 			try {
 				const auth = await awaitWithDeadline(
-					resolveUsageAuth(ctx, adapter),
+					resolveUsageAuth(ctx, adapter, undefined, credentialReader, credentialCandidates),
 					signal,
 					DEFAULT_TIMEOUT_MS,
 					`revalidating ${adapter.displayName} runtime auth`,
@@ -438,7 +440,7 @@ export default function usageExtension(
 		if (!adapter) return false;
 		try {
 			const auth = await awaitWithDeadline(
-				resolveUsageAuth(ctx, adapter),
+				resolveUsageAuth(ctx, adapter, undefined, credentialReader, credentialCandidates),
 				signal,
 				DEFAULT_TIMEOUT_MS,
 				`revalidating ${adapter.displayName} runtime auth`,
@@ -677,7 +679,7 @@ export default function usageExtension(
 								async (signal) => {
 									const expectedModel = modelIdentity(ctx.model);
 									const auth = await awaitWithDeadline(
-										resolveCodexResetAuth(ctx, undefined, credentialReader),
+										resolveCodexResetAuth(ctx, undefined, credentialReader, credentialCandidates),
 										signal,
 										DEFAULT_TIMEOUT_MS,
 										"resolving current Codex reset authentication",
@@ -695,7 +697,7 @@ export default function usageExtension(
 										};
 									}
 									const revalidated = await awaitWithDeadline(
-										resolveCodexResetAuth(ctx, undefined, credentialReader),
+										resolveCodexResetAuth(ctx, undefined, credentialReader, credentialCandidates),
 										signal,
 										DEFAULT_TIMEOUT_MS,
 										"revalidating current Codex reset authentication",
@@ -759,7 +761,7 @@ export default function usageExtension(
 								controller.signal,
 								async (signal) => {
 									const auth = await awaitWithDeadline(
-										resolveCodexResetAuth(ctx, undefined, credentialReader),
+										resolveCodexResetAuth(ctx, undefined, credentialReader, credentialCandidates),
 										signal,
 										DEFAULT_TIMEOUT_MS,
 										"revalidating current Codex reset authentication",