pi-multikey 1.2.0

package/package.json

@@ -0,0 +1,44 @@
+{
+	"name": "pi-multikey",
+	"version": "1.2.0",
+	"description": "One pi provider backed by many API keys: automatic 429 rotation, per-request key leases for concurrent subagents, and a /multikey management TUI",
+	"keywords": [
+		"pi-package",
+		"pi",
+		"pi-extension",
+		"multikey",
+		"api-keys",
+		"rate-limit",
+		"429",
+		"provider",
+		"key-pool"
+	],
+	"license": "MIT",
+	"repository": {
+		"type": "git",
+		"url": "git+https://github.com/kslamph/multikey.git"
+	},
+	"files": [
+		"index.ts",
+		"config.ts",
+		"presets.ts",
+		"probe.ts",
+		"pool.ts",
+		"stream.ts",
+		"manage.ts",
+		"tui.ts",
+		"README.md",
+		"README.zh.md",
+		"LICENSE"
+	],
+	"peerDependencies": {
+		"@earendil-works/pi-ai": "*",
+		"@earendil-works/pi-coding-agent": "*",
+		"@earendil-works/pi-tui": "*"
+	},
+	"pi": {
+		"extensions": [
+			"./index.ts"
+		]
+	}
+}

package/pool.ts

@@ -0,0 +1,180 @@
+/**
+ * KeyPool: per-pool API key state, lease acquisition, cooldowns, stats.
+ *
+ * Selection policy: among enabled keys not in cooldown, pick the one with the
+ * fewest in-flight requests, tie-broken by least-recently-used. This spreads
+ * concurrent subagents across different keys automatically.
+ */
+
+import type { PoolConfig } from "./config.ts";
+import { maskKey } from "./config.ts";
+
+export type KeyOutcome = "ok" | "rate_limited" | "invalid" | "error";
+
+export interface Lease {
+	key: string;
+	label: string;
+	acquiredAt: number;
+}
+
+interface KeyStat {
+	ok: number;
+	rateLimited: number;
+	invalid: number;
+	errors: number;
+	inflight: number;
+	cooldownUntil: number;
+	cooldownReason?: string;
+	lastUsed: number;
+}
+
+function newStat(): KeyStat {
+	return { ok: 0, rateLimited: 0, invalid: 0, errors: 0, inflight: 0, cooldownUntil: 0, lastUsed: 0 };
+}
+
+export class KeyPool {
+	config: PoolConfig;
+	private stats = new Map<string, KeyStat>();
+
+	constructor(config: PoolConfig) {
+		this.config = config;
+	}
+
+	private stat(key: string): KeyStat {
+		let s = this.stats.get(key);
+		if (!s) {
+			s = newStat();
+			this.stats.set(key, s);
+		}
+		return s;
+	}
+
+	private enabledKeys(): { key: string; label: string }[] {
+		return this.config.keys
+			.filter((k) => k.enabled !== false)
+			.map((k) => ({ key: k.key, label: k.label ?? maskKey(k.key) }));
+	}
+
+	get size(): number {
+		return this.enabledKeys().length;
+	}
+
+	mask(key: string): string {
+		return maskKey(key);
+	}
+
+	/** Replace config in place (keeps live stats). */
+	updateConfig(config: PoolConfig): void {
+		Object.assign(this.config, config);
+	}
+
+	/**
+	 * Acquire a key lease. If every key is in cooldown, waits (signal-aware)
+	 * until the earliest cooldown expires, then acquires that key.
+	 */
+	async acquire(signal?: AbortSignal): Promise<Lease> {
+		for (;;) {
+			if (signal?.aborted) throw new DOMException("Aborted", "AbortError");
+			const now = Date.now();
+			const candidates = this.enabledKeys()
+				.map((entry) => ({ ...entry, stat: this.stat(entry.key) }))
+				.filter((entry) => entry.stat.cooldownUntil <= now);
+
+			if (candidates.length > 0) {
+				candidates.sort((a, b) => a.stat.inflight - b.stat.inflight || a.stat.lastUsed - b.stat.lastUsed);
+				const chosen = candidates[0]!;
+				chosen.stat.inflight++;
+				chosen.stat.lastUsed = now;
+				return { key: chosen.key, label: chosen.label, acquiredAt: now };
+			}
+
+			// All keys cooling: wait for the earliest recovery (bounded to 60s).
+			const cooldowns = this.enabledKeys().map((entry) => this.stat(entry.key).cooldownUntil);
+			const earliest = Math.min(...cooldowns);
+			const waitMs = Math.max(250, Math.min(earliest - now, 60_000));
+			await this.sleep(waitMs, signal);
+		}
+	}
+
+	release(lease: Lease): void {
+		const s = this.stat(lease.key);
+		s.inflight = Math.max(0, s.inflight - 1);
+	}
+
+	report(lease: Lease, outcome: KeyOutcome, cooldownMs?: number): void {
+		const s = this.stat(lease.key);
+		const now = Date.now();
+		switch (outcome) {
+			case "ok":
+				s.ok++;
+				s.cooldownUntil = 0;
+				s.cooldownReason = undefined;
+				break;
+			case "rate_limited":
+				s.rateLimited++;
+				s.cooldownUntil = Math.max(s.cooldownUntil, now + (cooldownMs ?? this.config.cooldownMs ?? 20_000));
+				s.cooldownReason = "429";
+				break;
+			case "invalid":
+				s.invalid++;
+				s.cooldownUntil = Math.max(s.cooldownUntil, now + (this.config.invalidKeyCooldownMs ?? 600_000));
+				s.cooldownReason = "invalid";
+				break;
+			case "error":
+				s.errors++;
+				break;
+		}
+	}
+
+	private sleep(ms: number, signal?: AbortSignal): Promise<void> {
+		return new Promise((resolve, reject) => {
+			const timer = setTimeout(() => {
+				signal?.removeEventListener("abort", onAbort);
+				resolve();
+			}, ms);
+			const onAbort = () => {
+				clearTimeout(timer);
+				reject(new DOMException("Aborted", "AbortError"));
+			};
+			if (signal) {
+				if (signal.aborted) {
+					clearTimeout(timer);
+					reject(new DOMException("Aborted", "AbortError"));
+					return;
+				}
+				signal.addEventListener("abort", onAbort, { once: true });
+			}
+		});
+	}
+
+	/** One row per configured key, for TUI status display. */
+	statusRows(): {
+		label: string;
+		masked: string;
+		enabled: boolean;
+		inflight: number;
+		ok: number;
+		rateLimited: number;
+		invalid: number;
+		errors: number;
+		cooldownRemainingMs: number;
+		cooldownReason?: string;
+	}[] {
+		const now = Date.now();
+		return this.config.keys.map((k) => {
+			const s = this.stat(k.key);
+			return {
+				label: k.label ?? maskKey(k.key),
+				masked: maskKey(k.key),
+				enabled: k.enabled !== false,
+				inflight: s.inflight,
+				ok: s.ok,
+				rateLimited: s.rateLimited,
+				invalid: s.invalid,
+				errors: s.errors,
+				cooldownRemainingMs: Math.max(0, s.cooldownUntil - now),
+				cooldownReason: s.cooldownReason,
+			};
+		});
+	}
+}

package/presets.ts

@@ -0,0 +1,143 @@
+/**
+ * Built-in provider presets.
+ *
+ * A preset carries everything except credentials, so adding a provider is just
+ * "paste your keys". Model specs here are verified against provider docs and
+ * live probes — see README "Presets" for the evidence trail.
+ */
+
+import type { PoolConfig, PoolModelConfig } from "./config.ts";
+
+export interface Preset {
+	/** Stable id used in the config file's `_preset` marker. */
+	id: string;
+	name: string;
+	description: string;
+	/** Suggested pi provider id (user can override). */
+	defaultPoolId: string;
+	baseUrl: string;
+	api: string;
+	compat?: Record<string, unknown>;
+	/** Where to create/copy API keys. */
+	keyHint?: string;
+	models: PoolModelConfig[];
+}
+
+const LEVELS = ["off", "minimal", "low", "medium", "high", "xhigh", "max"] as const;
+
+/** Build a thinkingLevelMap that always writes every level (null = hidden in the UI). */
+function levels(
+	supported: Partial<Record<(typeof LEVELS)[number], string | null>>,
+): Record<string, string | null> {
+	const map: Record<string, string | null> = {};
+	for (const level of LEVELS) {
+		map[level] = level in supported ? (supported[level] ?? null) : null;
+	}
+	return map;
+}
+
+const BAI_COMPAT = {
+	supportsDeveloperRole: false,
+	thinkingFormat: "deepseek",
+	requiresReasoningContentOnAssistantMessages: true,
+	supportsReasoningEffort: true,
+};
+
+export const PRESETS: Preset[] = [
+	{
+		id: "b-ai",
+		name: "B.AI",
+		description: "api.b.ai — DeepSeek V4 Flash (+vision), Hunyuan Hy3, MiMo V2.5, Qwen3.8, GLM 5.3 (6 models)",
+		defaultPoolId: "bai",
+		baseUrl: "https://api.b.ai/v1",
+		api: "openai-completions",
+		compat: BAI_COMPAT,
+		keyHint: "https://www.b.ai/ → API Keys (one entry per key; multiple keys share the load)",
+		models: [
+			{
+				// docs.b.ai/llmservice/models/deepseek-v4-flash + api-docs.deepseek.com/guides/thinking_mode
+				// Native effort tiers: low/high/max (medium & xhigh alias to high server-side).
+				id: "deepseek-v4-flash",
+				name: "DeepSeek V4 Flash",
+				reasoning: true,
+				input: ["text"],
+				contextWindow: 1_000_000,
+				maxTokens: 384_000,
+				thinkingLevelMap: levels({ off: "none", low: "low", high: "high", max: "max" }),
+			},
+			{
+				// No public card; experimental vision build on the same V4 Flash backbone.
+				id: "deepseek-v4-flash-vision-exp",
+				name: "DeepSeek V4 Flash Vision (exp)",
+				reasoning: true,
+				input: ["text", "image"],
+				contextWindow: 1_000_000,
+				maxTokens: 384_000,
+				thinkingLevelMap: levels({ off: "none", low: "low", high: "high", max: "max" }),
+			},
+			{
+				// docs.b.ai/llmservice/models/hy3 — Tencent modes: no_think / think_low / think_high.
+				id: "hy3",
+				name: "Hunyuan Hy3",
+				reasoning: true,
+				input: ["text"],
+				contextWindow: 262_144,
+				maxTokens: 131_072,
+				thinkingLevelMap: levels({ off: "none", low: "low", high: "high" }),
+			},
+			{
+				// docs.b.ai/llmservice/models/mimo-v2.5 + Xiaomi official: only `none` disables
+				// thinking; low/medium/high are accepted but behave identically. b.ai rejects
+				// minimal/xhigh/max with HTTP 400.
+				id: "mimo-v2.5",
+				name: "Xiaomi MiMo V2.5",
+				reasoning: true,
+				input: ["text", "image"],
+				contextWindow: 1_000_000,
+				maxTokens: 131_072,
+				thinkingLevelMap: levels({ off: "none", high: "high" }),
+			},
+			{
+				// docs.b.ai/llmservice/models/qwen3-8-flash — tiers offered: off/low/medium/xhigh.
+				id: "qwen3.8-flash",
+				name: "Qwen3.8 Flash",
+				reasoning: true,
+				input: ["text", "image"],
+				contextWindow: 1_000_000,
+				maxTokens: 131_072,
+				thinkingLevelMap: levels({ off: "none", low: "low", medium: "medium", xhigh: "xhigh" }),
+			},
+			{
+				// docs.b.ai/llmservice/models/glm-5-3-flash — always thinks (off unsupported);
+				// reasoning_effort: low/high/max, default max.
+				id: "glm-5.3-flash",
+				name: "GLM 5.3 Flash",
+				reasoning: true,
+				input: ["text", "image"],
+				contextWindow: 1_000_000,
+				maxTokens: 131_072,
+				thinkingLevelMap: levels({ low: "low", high: "high", max: "max" }),
+			},
+		],
+	},
+];
+
+export function findPreset(id: string): Preset | undefined {
+	return PRESETS.find((p) => p.id === id);
+}
+
+/** Materialize a preset into a pool config with the given keys. */
+export function poolFromPreset(preset: Preset, poolId: string, keys: string[]): PoolConfig {
+	return {
+		id: poolId,
+		name: `${preset.name} (Key Pool)`,
+		baseUrl: preset.baseUrl,
+		api: preset.api,
+		compat: preset.compat ? { ...preset.compat } : undefined,
+		cooldownMs: 20_000,
+		invalidKeyCooldownMs: 600_000,
+		keys: keys.map((key, i) => ({ key, label: `key-${i + 1}`, enabled: true })),
+		// Deep copy so per-pool edits never mutate the shipped preset.
+		models: JSON.parse(JSON.stringify(preset.models)) as PoolModelConfig[],
+	};
+}

package/probe.ts

@@ -0,0 +1,264 @@
+/**
+ * Endpoint probing: auto-detect the auth header style and fetch the model list.
+ *
+ * Strategy (no questions asked):
+ *   1. GET <baseUrl>/models with `Authorization: Bearer <key>` — if 200, we have
+ *      the model list. If 401/403, retry with `x-api-key: <key>`.
+ *   2. /models is public on some gateways, so a 200 there doesn't prove the key
+ *      works. We therefore verify auth with a tiny 1-token chat request using
+ *      the detected style; 401/403 there rotates to the other style.
+ *   3. If even the chat probe accepts a bogus key, the endpoint simply doesn't
+ *      check keys — we proceed with the default (bearer) and say so.
+ *
+ * Everything degrades gracefully: probe failures never block pool creation,
+ * they only downgrade to manual model entry.
+ */
+
+export type AuthStyle = "bearer" | "api-key";
+
+export interface RemoteModel {
+	id: string;
+	name?: string;
+	contextWindow?: number;
+	maxTokens?: number;
+	input?: ("text" | "image")[];
+	reasoning?: boolean;
+}
+
+export interface ProbeResult {
+	/** True when a model list was fetched (even from an open endpoint). */
+	ok: boolean;
+	/** Auth header style that verified (or the safe default when unverifiable). */
+	auth: AuthStyle;
+	/** "confirmed" = chat probe proved it · "unverified" = endpoint open / unreachable · "rejected" = both styles got 401/403. */
+	authStatus: "confirmed" | "unverified" | "rejected";
+	modelsUrl?: string;
+	models?: RemoteModel[];
+	/** Human-readable attempt log for transparency in the TUI. */
+	log: string[];
+}
+
+const MODELS_TIMEOUT_MS = 12_000;
+const CHAT_TIMEOUT_MS = 25_000;
+
+function trimSlash(url: string): string {
+	return url.replace(/\/+$/, "");
+}
+
+/** Candidate /models URLs: <base>/models, and <base>/v1/models when the base doesn't already end in a version segment. */
+export function modelsUrlCandidates(baseUrl: string): string[] {
+	const base = trimSlash(baseUrl);
+	const candidates = [`${base}/models`];
+	if (!/\/v\d+(?:beta|preview)?$/.test(base)) candidates.push(`${base}/v1/models`);
+	return [...new Set(candidates)];
+}
+
+function authHeaders(style: AuthStyle, key: string): Record<string, string> {
+	return style === "bearer" ? { Authorization: `Bearer ${key}` } : { "x-api-key": key };
+}
+
+async function fetchJson(url: string, headers: Record<string, string>, timeoutMs: number): Promise<{ status: number; body?: unknown }> {
+	try {
+		const response = await fetch(url, {
+			method: "GET",
+			headers: { Accept: "application/json", ...headers },
+			signal: AbortSignal.timeout(timeoutMs),
+		});
+		let body: unknown;
+		try {
+			body = await response.json();
+		} catch {
+			// Non-JSON body (HTML error page etc.) — status is still meaningful.
+		}
+		return { status: response.status, body };
+	} catch (error) {
+		throw new Error(error instanceof Error ? error.message : String(error));
+	}
+}
+
+function asRecord(value: unknown): Record<string, unknown> | undefined {
+	return value && typeof value === "object" && !Array.isArray(value) ? (value as Record<string, unknown>) : undefined;
+}
+
+function firstNumber(source: Record<string, unknown>, keys: string[]): number | undefined {
+	for (const key of keys) {
+		const value = source[key];
+		if (typeof value === "number" && Number.isFinite(value) && value > 0) return value;
+	}
+	return undefined;
+}
+
+/** Extract model entries from the many /models response shapes (OpenAI `{data:[]}`, bare array, `{models:[]}`). */
+export function parseModelsResponse(body: unknown): RemoteModel[] {
+	const root = asRecord(body);
+	const list = Array.isArray(body) ? body : (root?.data ?? root?.models);
+	if (!Array.isArray(list)) return [];
+
+	const models: RemoteModel[] = [];
+	for (const entry of list) {
+		if (typeof entry === "string") {
+			if (entry.trim()) models.push({ id: entry.trim() });
+			continue;
+		}
+		const record = asRecord(entry);
+		if (!record) continue;
+		const id = typeof record.id === "string" && record.id.trim()
+			? record.id.trim()
+			: typeof record.model === "string" && record.model.trim()
+				? record.model.trim()
+				: typeof record.name === "string" && record.name.trim()
+					? record.name.trim()
+					: undefined;
+		if (!id) continue;
+
+		const model: RemoteModel = { id };
+		if (typeof record.name === "string" && record.name.trim() && record.name.trim() !== id) model.name = record.name.trim();
+		else if (typeof record.display_name === "string" && record.display_name.trim() && record.display_name.trim() !== id) model.name = record.display_name.trim();
+
+		const contextWindow = firstNumber(record, ["context_length", "context_window", "max_model_len", "max_context_length", "context_size"]);
+		if (contextWindow) model.contextWindow = contextWindow;
+
+		const maxTokens = firstNumber(record, ["max_output_tokens", "max_completion_tokens", "output_token_limit"]);
+		if (maxTokens) model.maxTokens = maxTokens;
+
+		const modalities = record.input_modalities ?? asRecord(record.architecture)?.input_modalities;
+		if (Array.isArray(modalities)) {
+			const input: ("text" | "image")[] = ["text"];
+			if (modalities.some((m) => typeof m === "string" && m.toLowerCase() === "image")) input.push("image");
+			model.input = input;
+		} else if (typeof modalities === "string" && modalities.toLowerCase() === "image") {
+			model.input = ["text", "image"];
+		}
+
+		const reasoningFlag = asRecord(record.reasoning)?.supported ?? record.reasoning_supported;
+		if (typeof reasoningFlag === "boolean") model.reasoning = reasoningFlag;
+
+		models.push(model);
+	}
+	return models;
+}
+
+/**
+ * Verify a key with a minimal chat completion (a few tokens at most). Returns
+ * "ok" when auth was accepted (2xx, or 4xx that clearly got past auth like a
+ * bad-model/params 400/404), "rejected" on 401/403, "error" on network trouble.
+ */
+async function chatProbe(baseUrl: string, style: AuthStyle, key: string, modelId: string): Promise<"ok" | "rejected" | "error"> {
+	try {
+		const response = await fetch(`${trimSlash(baseUrl)}/chat/completions`, {
+			method: "POST",
+			headers: { "Content-Type": "application/json", ...authHeaders(style, key) },
+			body: JSON.stringify({ model: modelId, max_tokens: 4, messages: [{ role: "user", content: "ping" }] }),
+			signal: AbortSignal.timeout(CHAT_TIMEOUT_MS),
+		});
+		if (response.status === 401 || response.status === 403) return "rejected";
+		return "ok"; // 2xx, or 4xx past auth (bad model / params) — auth itself worked.
+	} catch {
+		return "error";
+	}
+}
+
+/**
+ * Probe an endpoint with the user's key: detect auth style, verify the key,
+ * and fetch the selectable model list. `chatModelId` (from the fetched list or
+ * a previously configured model) enables the chat auth verification.
+ */
+export async function probeEndpoint(
+	baseUrl: string,
+	key: string,
+	options?: { chatModelId?: string; authHint?: AuthStyle; onLog?: (line: string) => void },
+): Promise<ProbeResult> {
+	const log: string[] = [];
+	const emit = (line: string) => {
+		log.push(line);
+		options?.onLog?.(line);
+	};
+	const styles: AuthStyle[] = options?.authHint === "api-key" ? ["api-key", "bearer"] : ["bearer", "api-key"];
+
+	// --- 1. Fetch the model list, rotating auth styles on 401/403. ---
+	let models: RemoteModel[] | undefined;
+	let modelsUrl: string | undefined;
+	let auth: AuthStyle = styles[0]!;
+	let listAuthRejected = true; // did /models itself reject every style?
+
+	for (const style of styles) {
+		for (const url of modelsUrlCandidates(baseUrl)) {
+			let result: { status: number; body?: unknown };
+			try {
+				emit(`GET ${url} (${style === "bearer" ? "Authorization: Bearer" : "x-api-key"})…`);
+				result = await fetchJson(url, authHeaders(style, key), MODELS_TIMEOUT_MS);
+			} catch (error) {
+				emit(`   network error: ${error instanceof Error ? error.message : String(error)}`);
+				continue;
+			}
+			if (result.status === 200) {
+				const parsed = parseModelsResponse(result.body);
+				if (parsed.length > 0) {
+					models = parsed;
+					modelsUrl = url;
+					auth = style;
+					listAuthRejected = false;
+					emit(`   200 — ${parsed.length} model(s) listed`);
+					break;
+				}
+				emit(`   200 but no models recognized in response`);
+				continue;
+			}
+			emit(`   HTTP ${result.status}`);
+			if (result.status !== 401 && result.status !== 403) continue; // 404 etc: try next URL form
+		}
+		if (models) break;
+	}
+
+	// --- 2. Verify auth with a tiny chat request (models lists can be public). ---
+	const chatModelId =
+		options?.chatModelId ?? (models && models.length > 0 ? models.find((m) => m.id.toLowerCase().includes("free"))?.id ?? models[0]!.id : undefined);
+	let authStatus: ProbeResult["authStatus"] = "unverified";
+
+	if (chatModelId) {
+		let verified: AuthStyle | undefined;
+		for (const style of styles) {
+			emit(`auth check: 1-token chat on "${chatModelId}" with ${style === "bearer" ? "Bearer" : "x-api-key"}…`);
+			const verdict = await chatProbe(baseUrl, style, key, chatModelId);
+			if (verdict === "ok") {
+				verified = style;
+				emit(`   accepted ✓ (style: ${style === "bearer" ? "Authorization: Bearer" : "x-api-key"})`);
+				break;
+			}
+			emit(`   ${verdict === "rejected" ? "rejected (401/403)" : "network error"}`);
+		}
+
+		if (verified) {
+			auth = verified;
+			authStatus = "confirmed";
+		} else if (listAuthRejected) {
+			// /models also rejected the key with every style — likely a bad key.
+			authStatus = "rejected";
+		} else {
+			// Models listed fine but chat auth unverifiable: maybe the endpoint is open.
+			try {
+				emit(`auth check: same chat with a dummy key to detect open endpoints…`);
+				const bogus = await chatProbe(baseUrl, "bearer", "multikey-open-endpoint-check", chatModelId);
+				if (bogus === "ok") {
+					emit(`   accepted — endpoint does not verify keys (open endpoint)`);
+				} else {
+					emit(`   rejected — key was not verified; double-check it`);
+					authStatus = "rejected";
+				}
+			} catch {
+				authStatus = "unverified";
+			}
+		}
+	} else {
+		emit("no model id available — skipping chat auth check");
+	}
+
+	return {
+		ok: models !== undefined,
+		auth,
+		authStatus,
+		modelsUrl,
+		models,
+		log,
+	};
+}