@adeildo/pi-providers 5.2.0 → 5.3.0

package/src/accounts/feature.ts

@@ -1,279 +1,141 @@
-import { defineFeature, literal, setting } from "@adeildo/pi-kit";
-import type { FeatureScope } from "@adeildo/pi-kit";
-import type { Credential, Provider } from "@earendil-works/pi-ai";
+// The accounts feature and its two entry points.
+import { ACCOUNT_STATE, type AccountState, defineFeature, isObject } from "@adeildo/pi-kit";
 import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
-import { Key } from "@earendil-works/pi-tui";
+import { type AutocompleteItem, Key } from "@earendil-works/pi-tui";
 
+import { claudeCodeVersion } from "../subscription/feature.ts";
+import { anthropicQuota } from "../usage/anthropic.ts";
+import { Quota } from "../usage/quota.ts";
+import { resetIn, windowsOf } from "../usage/read.ts";
+import { refreshStatus } from "./actions.ts";
 import { activeAccount } from "./active.ts";
-import { reason } from "./describe.ts";
-import { interactionFor } from "./interaction.ts";
-import { liftProvider, nativeOf, type AccountSession } from "./lift.ts";
-import { login, loginMethods, type LoginMethod } from "./login.ts";
-import { LOGIN_KEY, NAME, STATUS_KEY } from "./names.ts";
-import { pin, replay, type Pins } from "./pins.ts";
-import { accountRows, DEFAULT_ACCOUNT, DEFAULT_LABEL, refreshStatus } from "./screen.ts";
+import { attachProvider, knownReading, showPlans, type Watch } from "./attach.ts";
+import {
+	type Accounts,
+	accountCount,
+	loginProviders,
+	manageAccounts,
+	pickAccount,
+} from "./dialogs.ts";
+import { NAME, STATUS_KEY } from "./names.ts";
+import { replay, type Pins } from "./pins.ts";
+import { usageLimitLine } from "./policy.ts";
+import { accountRows } from "./screen.ts";
+import { onAuthFailure, onLimit } from "./settings.ts";
 import { AccountStore } from "./store.ts";
-import type { Account } from "./types.ts";
-
-/** What happens when an account hits the limit of its plan. */
-export const WHEN_LIMITED = ["ask", "switch", "stop"] as const;
-export type WhenLimited = (typeof WHEN_LIMITED)[number];
-
-export const onLimit = setting<WhenLimited>({
-	id: "accounts.onLimit",
-	default: "ask",
-	decoder: literal(...WHEN_LIMITED),
-	ui: {
-		section: "Accounts",
-		label: "When an account hits its limit",
-		description: "Ask me first, switch on its own, or stop and tell me.",
-		control: {
-			type: "choice",
-			options: [
-				{ value: "ask", label: "ask me first" },
-				{ value: "switch", label: "switch on its own" },
-				{ value: "stop", label: "stop and tell me" },
-			],
-		},
-	},
-});
 
 export const accounts = defineFeature({
 	id: "accounts",
 	description: "Use more than one credential of the same provider",
 	tab: "Providers",
-	settings: [onLimit],
+	settings: [onLimit, onAuthFailure],
 	setup(scope) {
 		const store = new AccountStore();
 		const pins: Pins = new Map();
+		const watch: Watch = {
+			limited: new Map(),
+			usage: new Map(),
+			quota: new Quota({
+				anthropic: anthropicQuota(() =>
+					scope.has("subscription") ? claudeCodeVersion.get(scope) : claudeCodeVersion.default,
+				),
+			}),
+		};
+
+		const bundle: Accounts = { scope, store, pins, watch };
+		let registry: ExtensionContext["modelRegistry"] | undefined;
 
 		scope.screen.rows((ctx) => accountRows(scope, store, pins, ctx));
-		scope.on("model_select", (_event, ctx) => refreshStatus(store, pins, ctx));
+		// A model of another provider brings that provider's plans, so its quota shows before the first
+		// answer instead of after it.
+		scope.on("model_select", (event, ctx) => {
+			refreshStatus(store, pins, ctx);
+			showPlans(store, pins, watch, ctx, event.model.provider);
+		});
 		scope.on("session_shutdown", (_event, ctx) => ctx.ui.setStatus(STATUS_KEY, undefined));
 
+		// Whoever draws the account in use asks for it on the spot, so it is never a frame old.
+		scope.events.on(ACCOUNT_STATE, (data: unknown) => {
+			if (!isObject(data) || typeof data.provider !== "string") return;
+			data.state = inUse(watch, store, pins, data.provider);
+		});
+
+		// The line stops pi's own retry.
+		scope.on("message_end", (event) => {
+			if (event.message.role !== "assistant") return;
+			const providerId = event.message.provider;
+			const account = activeAccount(store, pins, providerId);
+			if (account === undefined) return;
+			const note = resetIn(knownReading(watch, providerId, account.id));
+			const text = usageLimitLine(event.message, account.label, note);
+			if (text === undefined) return;
+			return { message: { ...event.message, errorMessage: text } };
+		});
+
 		scope.onSessionStart((ctx) => {
+			registry = ctx.modelRegistry;
 			for (const warning of store.reload()) scope.warn(warning);
 			pins.clear();
+			watch.limited.clear();
+			watch.usage.clear();
 			for (const [provider, account] of replay(ctx.sessionManager.getBranch())) {
 				pins.set(provider, account);
 			}
 			for (const providerId of new Set([...store.providerIds(), ...pins.keys()])) {
-				lift(scope, store, pins, ctx, providerId);
+				attachProvider(scope, store, pins, watch, ctx, providerId);
 			}
 			refreshStatus(store, pins, ctx);
+			showPlans(store, pins, watch, ctx, ctx.model?.provider);
 		});
 
 		scope.registerShortcut(Key.alt("a"), {
-			description: `${NAME}: pick the account of the current provider`,
-			handler: (ctx) => pickAccount(scope, store, pins, ctx),
+			description: `${NAME}: switch the account of the current provider for this session`,
+			handler: (ctx) => pickAccount(bundle, ctx),
 		});
 
 		scope.registerCommand("accounts", {
-			description: `${NAME}: add an account for a provider`,
-			handler: (args, ctx) => addAccount(scope, store, pins, args, ctx),
+			description: `${NAME}: see, switch, rename, remove or add the accounts of a provider`,
+			getArgumentCompletions: (prefix) => providerCompletions(registry, store, prefix),
+			handler: (args, ctx) => manageAccounts(bundle, args, ctx),
 		});
 	},
 });
 
-function sessionFor(
-	scope: FeatureScope,
+/** The providers `/accounts` can take, those with accounts first. */
+function providerCompletions(
+	registry: ExtensionContext["modelRegistry"] | undefined,
 	store: AccountStore,
-	pins: Pins,
-	ctx: ExtensionContext,
-	providerId: string,
-): AccountSession {
-	return {
-		resolve: () => {
-			const account = activeAccount(store, pins, providerId);
-			return account === undefined ? undefined : { id: account.id, credential: account.credential };
-		},
-		save: (id, credential) => void store.setCredential(providerId, id, credential),
-		afterLimit: (currentId, detail) =>
-			afterLimit(scope, store, pins, ctx, providerId, currentId, detail),
-	};
-}
-
-/** Follows the `accounts.onLimit` policy: ask, switch, or stop, and pins the account it moves to. */
-export async function afterLimit(
-	scope: FeatureScope,
-	store: AccountStore,
-	pins: Pins,
-	ctx: ExtensionContext,
-	providerId: string,
-	currentId: string,
-	detail: string,
-): Promise<{ id: string; credential: Credential } | undefined> {
-	const mode = onLimit.get(scope);
-	if (mode === "stop") return undefined;
-
-	const next = nextAccount(store.accounts(providerId), currentId);
-	if (next === undefined) return undefined;
-
-	if (mode === "ask") {
-		const once = `Switch to ${next.label}`;
-		const always = "Always switch when an account hits its limit";
-		const picked = await ctx.ui.select(`"${currentId}" hit its limit`, [once, always, "Stop here"]);
-		if (picked === always) scope.settings.set(onLimit, "switch");
-		else if (picked !== once) return undefined;
-	}
-
-	pin(scope, pins, providerId, next.id);
-	refreshStatus(store, pins, ctx);
-	ctx.ui.notify(`${NAME}: ${providerId} now uses ${next.label} (${detail})`, "warning");
-	return { id: next.id, credential: next.credential };
-}
-
-/** The account after this one in the list, so a limit moves on instead of repeating. */
-function nextAccount(list: readonly Account[], currentId: string): Account | undefined {
-	if (list.length < 2) return undefined;
-	const index = list.findIndex((account) => account.id === currentId);
-	return list[(index + 1) % list.length];
+	prefix: string,
+): AutocompleteItem[] | null {
+	if (registry === undefined) return null;
+	const typed = prefix.trim().toLowerCase();
+	const items = loginProviders(registry, store)
+		.filter((id) => id.toLowerCase().startsWith(typed))
+		.map((id) => {
+			const name = registry.getProviderDisplayName(id);
+			const count = store.accounts(id).length;
+			return {
+				value: id,
+				label: id,
+				description: count === 0 ? name : `${name}, ${accountCount(count)}`,
+			};
+		});
+	return items.length === 0 ? null : items;
 }
 
-function lift(
-	scope: FeatureScope,
+/** The account in use for a provider, with the windows it is spending. */
+function inUse(
+	watch: Watch,
 	store: AccountStore,
 	pins: Pins,
-	ctx: ExtensionContext,
 	providerId: string,
-): void {
-	// A pin with no account left behind it injects nothing, so re-registering it would only cost a
-	// model-registry rebuild.
-	if (!store.has(providerId)) return;
-
-	// `getProvider` returns the lifted provider once we registered one, so unwrap it before wrapping
-	// again. After a `/reload` it returns the built-in one, which is unwrapped too.
-	const provider = ctx.modelRegistry.getProvider(providerId);
-	if (provider === undefined) {
-		scope.warn(`no provider "${providerId}", so its accounts do nothing`);
-		return;
-	}
-	scope.registerProvider(
-		liftProvider(nativeOf(provider), sessionFor(scope, store, pins, ctx, providerId)),
-	);
-}
-
-async function pickAccount(
-	scope: FeatureScope,
-	store: AccountStore,
-	pins: Pins,
-	ctx: ExtensionContext,
-): Promise<void> {
-	if (!ctx.hasUI) {
-		scope.warn("no dialog to pick an account");
-		return;
-	}
-
-	const providerId = ctx.model?.provider;
-	if (providerId === undefined) {
-		ctx.ui.notify(`${NAME}: no model selected`, "warning");
-		return;
-	}
-	const entries = store.accounts(providerId);
-	if (entries.length === 0) {
-		ctx.ui.notify(`${NAME}: ${providerId} has no extra account yet`, "info");
-		return;
-	}
-
-	const current = activeAccount(store, pins, providerId)?.id ?? DEFAULT_ACCOUNT;
-	const choices = [
-		{ id: DEFAULT_ACCOUNT, label: DEFAULT_LABEL, kind: "login" },
-		...entries.map((account) => ({
-			id: account.id,
-			label: account.label,
-			kind: account.credential.type === "oauth" ? "oauth" : "key",
-		})),
-	];
-	const options = choices.map((choice, index) => {
-		const here = choice.id === current ? " ✓" : "";
-		return `${index + 1}  ${choice.label} (${choice.kind})${here}`;
-	});
-
-	const picked = await ctx.ui.select(`Account for ${providerId}`, options);
-	const index = picked === undefined ? -1 : options.indexOf(picked);
-	if (index < 0) return;
-
-	const choice = choices[index];
-	if (choice === undefined) return;
-	pin(scope, pins, providerId, choice.id === DEFAULT_ACCOUNT ? null : choice.id);
-	refreshStatus(store, pins, ctx);
-	ctx.ui.notify(`${NAME}: ${providerId} uses ${choice.label}`, "info");
-}
-
-async function addAccount(
-	scope: FeatureScope,
-	store: AccountStore,
-	pins: Pins,
-	args: string,
-	ctx: ExtensionContext,
-): Promise<void> {
-	if (!ctx.hasUI) {
-		scope.warn("no dialog to add an account");
-		return;
-	}
-
-	const providerId = args.trim() || (await pickProvider(ctx));
-	if (providerId === undefined) return;
-	const provider = ctx.modelRegistry.getProvider(providerId);
-	if (provider === undefined) {
-		ctx.ui.notify(`${NAME}: no provider "${providerId}"`, "error");
-		return;
-	}
-	const method = await pickMethod(ctx, provider);
-	if (method === undefined) return;
-
-	const label =
-		(await ctx.ui.input(`Name this ${providerId} account`, providerId))?.trim() || providerId;
-	try {
-		const credential = await login(
-			method,
-			interactionFor(ctx, providerId, new AbortController().signal),
-		);
-		const added = store.add(providerId, label, credential);
-		if (added.account === undefined) {
-			ctx.ui.notify(`${NAME}: ${added.problem}`, "error");
-			return;
-		}
-		store.setActive(providerId, added.account.id);
-		pin(scope, pins, providerId, added.account.id);
-		// The first account of a provider was not lifted at session start, so lift it now.
-		if (store.accounts(providerId).length === 1) lift(scope, store, pins, ctx, providerId);
-		refreshStatus(store, pins, ctx);
-		ctx.ui.notify(`${NAME}: ${providerId} account "${label}" added`, "info");
-	} catch (error) {
-		ctx.ui.notify(`${NAME}: ${reason(error)}`, "error");
-	} finally {
-		ctx.ui.setStatus(LOGIN_KEY, undefined);
-	}
-}
-
-async function pickProvider(ctx: ExtensionContext): Promise<string | undefined> {
-	const candidates: string[] = [];
-	for (const providerId of new Set(ctx.modelRegistry.getAll().map((model) => model.provider))) {
-		const provider = ctx.modelRegistry.getProvider(providerId);
-		if (provider !== undefined && loginMethods(provider).length > 0) candidates.push(providerId);
-	}
-	const list = candidates.toSorted();
-	if (list.length === 0) {
-		ctx.ui.notify(`${NAME}: no provider offers a login`, "warning");
-		return undefined;
-	}
-	if (list.length === 1) return list[0];
-	return await ctx.ui.select("Provider", list);
-}
-
-async function pickMethod(
-	ctx: ExtensionContext,
-	provider: Provider,
-): Promise<LoginMethod | undefined> {
-	const methods = loginMethods(provider);
-	if (methods.length === 0) {
-		ctx.ui.notify(`${NAME}: ${provider.id} has no login method`, "warning");
-		return undefined;
-	}
-	if (methods.length === 1) return methods[0];
-
-	const labels = methods.map((method) => method.label);
-	const picked = await ctx.ui.select(`${provider.id} login`, labels);
-	return methods.find((method) => method.label === picked);
+): AccountState | undefined {
+	const account = activeAccount(store, pins, providerId);
+	if (account === undefined) return undefined;
+	return {
+		provider: providerId,
+		label: account.label,
+		windows: windowsOf(knownReading(watch, providerId, account.id)),
+		...(account.health?.lastError?.kind === "auth" ? { needsLogin: true } : {}),
+	};
 }

package/src/accounts/interaction.ts

@@ -1,11 +1,23 @@
-// The provider's own login asks for a code, a URL or a choice. These dialogs answer for it, so no
-// OAuth and no key prompt live in this package.
+// The provider's own login asks for a code, a URL or a choice. Our framed screen answers it; a host
+// that cannot draw one falls back to pi's own dialogs.
 import type { AuthEvent, ProviderAuthInteraction } from "@earendil-works/pi-ai";
 import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
 
+import { type LoginSession, loginScreen } from "../ui/login.ts";
+import { pick } from "../ui/picker.ts";
 import { LOGIN_KEY } from "./names.ts";
 
-export function interactionFor(
+export async function loginFor(
+	ctx: ExtensionContext,
+	label: string,
+	signal: AbortSignal,
+): Promise<LoginSession> {
+	const screen = await loginScreen(ctx, label, signal);
+	return screen ?? { interaction: dialogs(ctx, label, signal), close: () => {} };
+}
+
+/** Pi's own dialogs, for a host without our screen. */
+function dialogs(
 	ctx: ExtensionContext,
 	label: string,
 	signal: AbortSignal,
@@ -15,7 +27,7 @@ export function interactionFor(
 		async prompt(prompt) {
 			if (prompt.type === "select") {
 				const options = prompt.options.map((option) => option.label);
-				const picked = await ctx.ui.select(prompt.message, options, { signal: prompt.signal });
+				const picked = await pick(ctx, prompt.message, options, { signal: prompt.signal });
 				if (picked === undefined) throw new Error(`${label}: login cancelled`);
 				return prompt.options.find((option) => option.label === picked)?.id ?? picked;
 			}