pi-multikey 1.2.0

package/stream.ts

@@ -0,0 +1,237 @@
+/**
+ * Rotating streamSimple: wraps the underlying pi-ai API implementation, picks a
+ * key from the pool for every request, and transparently retries on 429/401/403
+ * with the next key (marking a cooldown on the failed key).
+ *
+ * Events are only relayed to the caller once the attempt is known to be healthy,
+ * so a rotated attempt never produces duplicate/partial output.
+ */
+
+import {
+	type AssistantMessage,
+	type AssistantMessageEvent,
+	type AssistantMessageEventStream,
+	type Api,
+	type Context,
+	createAssistantMessageEventStream,
+	getApiProvider,
+	type Model,
+	type SimpleStreamOptions,
+} from "@earendil-works/pi-ai";
+import type { KeyOutcome, KeyPool, Lease } from "./pool.ts";
+
+const RATE_LIMIT_RE = /\b429\b|rate\s*limit|too many requests|quota\s*(exceed|limit)|requests per minute|requests per day/i;
+const INVALID_KEY_RE = /\b40[13]\b|unauthorized|forbidden|invalid\s*(api\s*)?key|incorrect\s*(api\s*)?key|authentication/i;
+
+interface CapturedResponse {
+	status: number;
+	retryAfterMs?: number;
+}
+
+export type Notifier = (message: string) => void;
+
+export function createRotatingStreamSimple(pool: KeyPool, apiName: string, notify: Notifier) {
+	const impl = getApiProvider(apiName as Api);
+	if (!impl) throw new Error(`multikey: no API provider registered for api: ${apiName}`);
+	// Auth style: "api-key" providers want the key in x-api-key (some reject
+	// Authorization entirely); bearer is the pi-ai default and needs no help.
+	const authStyle = pool.config.auth ?? "bearer";
+
+	return function rotatingStreamSimple(
+		model: Model<Api>,
+		context: Context,
+		options?: SimpleStreamOptions,
+	): AssistantMessageEventStream {
+		const out = createAssistantMessageEventStream();
+
+		void (async () => {
+			const maxAttempts = Math.max(1, pool.size);
+			let lastProblem = "no attempts made";
+
+			for (let attempt = 0; attempt < maxAttempts; attempt++) {
+				let lease: Lease;
+				try {
+					lease = await pool.acquire(options?.signal);
+				} catch {
+					emitAborted(out, model);
+					return;
+				}
+
+				try {
+					const captured: CapturedResponse = { status: 0 };
+					const attemptOptions: SimpleStreamOptions = {
+						...options,
+						apiKey: lease.key,
+						headers: authStyle === "api-key" ? { ...options?.headers, "x-api-key": lease.key } : options?.headers,
+						onResponse: (response) => {
+							captured.status = response.status;
+							const ra = response.headers?.["retry-after"];
+							const seconds = typeof ra === "string" ? Number(ra) : Number.NaN;
+							if (Number.isFinite(seconds) && seconds > 0) captured.retryAfterMs = seconds * 1000;
+						},
+					};
+
+					const sub = impl.streamSimple(model, context, attemptOptions);
+					const verdict = await pump(sub, out, captured);
+
+					if (verdict.kind === "completed") {
+						pool.report(lease, "ok");
+						return;
+					}
+
+					// Attempt failed. Only rotate when nothing was relayed yet; if content
+					// already streamed, the failure is surfaced as-is (mid-stream 429 is rare).
+					if (verdict.kind === "rotate" && !verdict.relayedAny) {
+						lastProblem = verdict.problem;
+						pool.report(lease, verdict.outcome, verdict.outcome === "rate_limited" ? verdict.cooldownMs : undefined);
+						notify(
+							`multikey[${pool.config.id}]: ${verdict.outcome === "rate_limited" ? "429" : "auth error"} on ${lease.label} (${pool.mask(lease.key)}), rotating to another key`,
+						);
+						continue;
+					}
+
+					// Non-rotatable failure relayed to caller already.
+					pool.report(lease, verdict.kind === "rotate" ? verdict.outcome : "error");
+					return;
+				} catch (error) {
+					if (options?.signal?.aborted) {
+						emitAborted(out, model);
+						return;
+					}
+					pool.report(lease, "error");
+					lastProblem = error instanceof Error ? error.message : String(error);
+					notify(`multikey[${pool.config.id}]: request error on ${lease.label}, trying another key`);
+				} finally {
+					pool.release(lease);
+				}
+			}
+
+			// All keys exhausted — surface a rate-limit-flavored error so pi's own
+			// retry/backoff kicks in; by then some cooldowns have expired.
+			emitError(out, model, `multikey[${pool.config.id}]: all ${maxAttempts} keys exhausted (last: ${lastProblem})`);
+		})();
+
+		return out;
+	};
+}
+
+type Verdict =
+	| { kind: "completed" }
+	| { kind: "rotate"; outcome: KeyOutcome; problem: string; cooldownMs?: number; relayedAny: boolean }
+	| { kind: "failed"; relayedAny: boolean };
+
+/**
+ * Relay events from `sub` into `out`. Buffer the initial "start" event until we
+ * know the HTTP status, so a 429 attempt can be dropped without the caller ever
+ * seeing it.
+ */
+async function pump(
+	sub: AsyncIterable<AssistantMessageEvent>,
+	out: AssistantMessageEventStream,
+	captured: CapturedResponse,
+): Promise<Verdict> {
+	const buffered: AssistantMessageEvent[] = [];
+	let relayedAny = false;
+
+	for await (const event of sub) {
+		if (event.type === "done") {
+			flush(buffered, out);
+			out.push(event);
+			out.end();
+			return { kind: "completed" };
+		}
+
+		if (event.type === "error") {
+			const problem = event.error?.errorMessage ?? "unknown provider error";
+			const rotation = classify(problem, captured);
+			if (rotation && !relayedAny) {
+				return { kind: "rotate", outcome: rotation.outcome, problem, cooldownMs: rotation.cooldownMs, relayedAny };
+			}
+			flush(buffered, out);
+			out.push(event);
+			out.end();
+			return { kind: "failed", relayedAny };
+		}
+
+		if (!relayedAny) {
+			if (captured.status !== 0) {
+				const rotation = classifyStatus(captured);
+				if (rotation) {
+					return { kind: "rotate", outcome: rotation.outcome, problem: `HTTP ${captured.status}`, cooldownMs: rotation.cooldownMs, relayedAny };
+				}
+				flush(buffered, out);
+				out.push(event);
+				relayedAny = true;
+				continue;
+			}
+			// Response headers not seen yet: buffer (this can only be "start").
+			buffered.push(event);
+			continue;
+		}
+
+		out.push(event);
+	}
+
+	// Stream ended without done/error (should not happen) — treat as failure.
+	if (!relayedAny && buffered.length === 0) {
+		return { kind: "rotate", outcome: "error", problem: "stream ended without events", relayedAny };
+	}
+	flush(buffered, out);
+	out.end();
+	return { kind: "failed", relayedAny: true };
+}
+
+function flush(buffered: AssistantMessageEvent[], out: AssistantMessageEventStream): void {
+	for (const event of buffered) out.push(event);
+	buffered.length = 0;
+}
+
+function classifyStatus(
+	captured: CapturedResponse,
+): { outcome: KeyOutcome; cooldownMs?: number } | undefined {
+	if (captured.status === 429) return { outcome: "rate_limited", cooldownMs: captured.retryAfterMs };
+	if (captured.status === 401 || captured.status === 403) return { outcome: "invalid" };
+	return undefined;
+}
+
+function classify(
+	message: string,
+	captured: CapturedResponse,
+): { outcome: KeyOutcome; cooldownMs?: number } | undefined {
+	const fromStatus = classifyStatus(captured);
+	if (fromStatus) return fromStatus;
+	if (RATE_LIMIT_RE.test(message)) return { outcome: "rate_limited" };
+	if (INVALID_KEY_RE.test(message)) return { outcome: "invalid" };
+	return undefined;
+}
+
+function baseMessage(model: Model<Api>, stopReason: "error" | "aborted", errorMessage?: string): AssistantMessage {
+	return {
+		role: "assistant",
+		content: [],
+		api: model.api,
+		provider: model.provider,
+		model: model.id,
+		usage: {
+			input: 0,
+			output: 0,
+			cacheRead: 0,
+			cacheWrite: 0,
+			totalTokens: 0,
+			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+		},
+		stopReason,
+		errorMessage,
+		timestamp: Date.now(),
+	};
+}
+
+function emitError(out: AssistantMessageEventStream, model: Model<Api>, message: string): void {
+	out.push({ type: "error", reason: "error", error: baseMessage(model, "error", message) });
+	out.end();
+}
+
+function emitAborted(out: AssistantMessageEventStream, model: Model<Api>): void {
+	out.push({ type: "error", reason: "aborted", error: baseMessage(model, "aborted") });
+	out.end();
+}

package/tui.ts

@@ -0,0 +1,349 @@
+/**
+ * TUI helpers in the npm:better-custom style: searchable single-select list,
+ * read-only info panel, built on ctx.ui.custom + @earendil-works/pi-tui.
+ */
+
+import { Key, matchesKey, truncateToWidth } from "@earendil-works/pi-tui";
+
+export interface SelectItem {
+	value: string;
+	label: string;
+	suffix?: string;
+	description?: string;
+}
+
+type CommandContext = Parameters<Parameters<import("@earendil-works/pi-coding-agent").ExtensionAPI["registerCommand"]>[1]["handler"]>[1];
+
+export async function selectOne(
+	ctx: CommandContext,
+	title: string,
+	items: SelectItem[],
+	options?: { initialIndex?: number },
+): Promise<string | null> {
+	if (items.length === 0) return null;
+
+	return await ctx.ui.custom<string | null>((tui, theme, _kb, done) => {
+		let cursor = Math.max(0, Math.min(options?.initialIndex ?? 0, items.length - 1));
+		let query = "";
+		let cachedLines: string[] | undefined;
+		const maxVisible = 12;
+
+		function visible(): SelectItem[] {
+			const q = query.trim().toLowerCase();
+			if (!q) return items;
+			return items.filter((item) =>
+				`${item.label} ${item.suffix ?? ""} ${item.description ?? ""}`.toLowerCase().includes(q),
+			);
+		}
+
+		function refresh() {
+			const v = visible();
+			if (v.length === 0) cursor = 0;
+			else if (cursor >= v.length) cursor = v.length - 1;
+			cachedLines = undefined;
+			tui.requestRender();
+		}
+
+		return {
+			render(width: number) {
+				if (cachedLines) return cachedLines;
+				const v = visible();
+				const w = Math.max(10, width);
+				const lines: string[] = [];
+				const add = (line = "") => lines.push(truncateToWidth(line, w));
+				const border = theme.fg("accent", "─".repeat(w));
+
+				add(border);
+				add(` ${theme.fg("accent", theme.bold(title))}`);
+				add(` ${theme.fg("text", `Search: ${query || "-"}`)}`);
+				add();
+
+				if (v.length === 0) {
+					add(theme.fg("warning", " No matches."));
+				} else {
+					const start = Math.max(0, Math.min(cursor - Math.floor(maxVisible / 2), Math.max(0, v.length - maxVisible)));
+					const end = Math.min(v.length, start + maxVisible);
+					for (let i = start; i < end; i++) {
+						const item = v[i]!;
+						const active = i === cursor;
+						const prefix = active ? theme.fg("accent", "> ") : "  ";
+						const label = active ? theme.fg("accent", item.label) : theme.fg("text", item.label);
+						const suffix = item.suffix ? theme.fg("dim", item.suffix) : "";
+						add(`${prefix}${label}${suffix}`);
+						if (item.description) {
+							for (const line of item.description.split("\n")) {
+								add(`   ${theme.fg("muted", line)}`);
+							}
+						}
+					}
+					if (v.length > maxVisible) {
+						add();
+						add(theme.fg("dim", ` ${start + 1}-${end} of ${v.length}`));
+					}
+				}
+
+				add();
+				add(theme.fg("dim", " Type to search • ↑↓ move • enter confirm • backspace delete • esc cancel"));
+				add(border);
+
+				cachedLines = lines;
+				return lines;
+			},
+			invalidate() {
+				cachedLines = undefined;
+			},
+			handleInput(data: string) {
+				const v = visible();
+				if (matchesKey(data, Key.up)) {
+					if (v.length === 0) return;
+					cursor = cursor === 0 ? v.length - 1 : cursor - 1;
+					refresh();
+					return;
+				}
+				if (matchesKey(data, Key.down)) {
+					if (v.length === 0) return;
+					cursor = cursor === v.length - 1 ? 0 : cursor + 1;
+					refresh();
+					return;
+				}
+				if (matchesKey(data, Key.enter)) {
+					done(v[cursor]?.value ?? null);
+					return;
+				}
+				if (matchesKey(data, Key.escape)) {
+					done(null);
+					return;
+				}
+				if (data === "\u007f" || data === "\b") {
+					if (query.length > 0) {
+						query = query.slice(0, -1);
+						refresh();
+					}
+					return;
+				}
+				if (data >= " " && data !== "\u001b" && data !== "\r" && data !== "\n") {
+					query += data;
+					cursor = 0;
+					refresh();
+				}
+			},
+		};
+	});
+}
+
+export async function showInfo(ctx: CommandContext, title: string, lines: string[]): Promise<void> {
+	await ctx.ui.custom<void>((tui, theme, _kb, done) => {
+		let cachedLines: string[] | undefined;
+		return {
+			render(width: number) {
+				if (cachedLines) return cachedLines;
+				const w = Math.max(10, width);
+				const out: string[] = [];
+				const add = (line = "") => out.push(truncateToWidth(line, w));
+				const border = theme.fg("accent", "─".repeat(w));
+				add(border);
+				add(` ${theme.fg("accent", theme.bold(title))}`);
+				add();
+				for (const line of lines) add(` ${line}`);
+				add();
+				add(theme.fg("dim", " esc/enter close"));
+				add(border);
+				cachedLines = out;
+				return out;
+			},
+			invalidate() {
+				cachedLines = undefined;
+			},
+			handleInput(data: string) {
+				if (matchesKey(data, Key.escape) || matchesKey(data, Key.enter)) done();
+			},
+		};
+	});
+}
+
+/** Simple toggle list (multi-select), returns selected values or null on cancel. */
+export async function pickMany(
+	ctx: CommandContext,
+	title: string,
+	items: SelectItem[],
+	options?: { preselected?: string[] },
+): Promise<string[] | null> {
+	return await ctx.ui.custom<string[] | null>((tui, theme, _kb, done) => {
+		let cursor = 0;
+		let query = "";
+		const selected = new Set<string>(options?.preselected ?? []);
+		let cachedLines: string[] | undefined;
+		const maxVisible = 12;
+
+		function visible(): SelectItem[] {
+			const q = query.trim().toLowerCase();
+			if (!q) return items;
+			return items.filter((item) => `${item.label} ${item.description ?? ""}`.toLowerCase().includes(q));
+		}
+
+		function refresh() {
+			const v = visible();
+			if (v.length === 0) cursor = 0;
+			else if (cursor >= v.length) cursor = v.length - 1;
+			cachedLines = undefined;
+			tui.requestRender();
+		}
+
+		return {
+			render(width: number) {
+				if (cachedLines) return cachedLines;
+				const v = visible();
+				const w = Math.max(10, width);
+				const lines: string[] = [];
+				const add = (line = "") => lines.push(truncateToWidth(line, w));
+				const border = theme.fg("accent", "─".repeat(w));
+				add(border);
+				add(` ${theme.fg("accent", theme.bold(title))}`);
+				add(` ${theme.fg("text", `Search: ${query || "-"}`)}`);
+				add();
+				if (v.length === 0) {
+					add(theme.fg("warning", " No matches."));
+				} else {
+					for (let i = 0; i < Math.min(v.length, maxVisible); i++) {
+						const item = v[i]!;
+						const active = i === cursor;
+						const check = selected.has(item.value) ? theme.fg("accent", "✓") : " ";
+						const prefix = active ? theme.fg("accent", "> ") : "  ";
+						add(`${prefix}[${check}] ${active ? theme.fg("accent", item.label) : item.label}`);
+					}
+				}
+				add();
+				add(theme.fg("dim", " space toggle • enter confirm • esc cancel"));
+				add(border);
+				cachedLines = lines;
+				return lines;
+			},
+			invalidate() {
+				cachedLines = undefined;
+			},
+			handleInput(data: string) {
+				const v = visible();
+				if (matchesKey(data, Key.up)) {
+					if (v.length === 0) return;
+					cursor = cursor === 0 ? v.length - 1 : cursor - 1;
+					refresh();
+					return;
+				}
+				if (matchesKey(data, Key.down)) {
+					if (v.length === 0) return;
+					cursor = cursor === v.length - 1 ? 0 : cursor + 1;
+					refresh();
+					return;
+				}
+				if (data === " ") {
+					const item = v[cursor];
+					if (item) {
+						if (selected.has(item.value)) selected.delete(item.value);
+						else selected.add(item.value);
+						refresh();
+					}
+					return;
+				}
+				if (matchesKey(data, Key.enter)) {
+					done([...selected]);
+					return;
+				}
+				if (matchesKey(data, Key.escape)) {
+					done(null);
+					return;
+				}
+				if (data === "\u007f" || data === "\b") {
+					if (query.length > 0) {
+						query = query.slice(0, -1);
+						refresh();
+					}
+					return;
+				}
+				if (data >= " " && data !== "\u001b" && data !== "\r" && data !== "\n") {
+					query += data;
+					cursor = 0;
+					refresh();
+				}
+			},
+		};
+	});
+}
+
+export async function inputNumber(ctx: CommandContext, title: string, current: number): Promise<number | undefined> {
+	const raw = await ctx.ui.input(title, String(current));
+	if (raw === undefined || raw.trim() === "") return undefined;
+	const value = Number(raw.trim());
+	if (!Number.isFinite(value) || value <= 0) return undefined;
+	return value;
+}
+
+const SPINNER_FRAMES = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"];
+
+/**
+ * Run an async task behind a live status panel: `update(line)` appends progress
+ * lines while the task runs; the panel closes itself with the task's result.
+ */
+export async function withProgress<T>(
+	ctx: CommandContext,
+	title: string,
+	task: (update: (line: string) => void) => Promise<T>,
+): Promise<T> {
+	return await ctx.ui.custom<T>((tui, theme, _kb, done) => {
+		const lines: string[] = [];
+		let finished = false;
+		let cachedLines: string[] | undefined;
+		let frame = 0;
+		const spinner = setInterval(() => {
+			if (finished) return;
+			frame = (frame + 1) % SPINNER_FRAMES.length;
+			tui.requestRender();
+		}, 90);
+
+		function update(line: string) {
+			lines.push(line);
+			cachedLines = undefined;
+			tui.requestRender();
+		}
+
+		void (async () => {
+			try {
+				const result = await task(update);
+				finished = true;
+				clearInterval(spinner);
+				done(result);
+			} catch (error) {
+				finished = true;
+				clearInterval(spinner);
+				update(`error: ${error instanceof Error ? error.message : String(error)}`);
+				// Give the user a moment to see the error before the panel closes.
+				setTimeout(() => done(undefined as T), 2500);
+			}
+		})();
+
+		return {
+			render(width: number) {
+				if (cachedLines) return cachedLines;
+				const w = Math.max(10, width);
+				const out: string[] = [];
+				const add = (line = "") => out.push(truncateToWidth(line, w));
+				const border = theme.fg("accent", "─".repeat(w));
+				add(border);
+				add(` ${theme.fg("accent", theme.bold(title))} ${finished ? "" : theme.fg("accent", SPINNER_FRAMES[frame]!)}`);
+				add();
+				for (const line of lines) add(` ${theme.fg("dim", line)}`);
+				add();
+				add(theme.fg("dim", " working — please wait…"));
+				add(border);
+				cachedLines = out;
+				return out;
+			},
+			invalidate() {
+				cachedLines = undefined;
+			},
+			handleInput() {
+				// Input intentionally ignored while the task runs.
+				if (finished) done(undefined as T);
+			},
+		};
+	});
+}