@hu3rror/pi-failover 0.2.0

package/src/failover-engine.ts

@@ -0,0 +1,314 @@
+export type FailureKind =
+	| "unauthorized"
+	| "rate-limited"
+	| "overloaded"
+	| "provider-error"
+	| "network";
+
+export interface FailureObservation {
+	status?: number;
+	kind?: FailureKind;
+	retryAfterMs?: number;
+}
+
+export type KeySlot = "primary" | "backup" | `backup-${number}`;
+
+export interface FailoverAttempt {
+	providerId: string;
+	model: string;
+	keySlot: KeySlot;
+}
+
+export interface KeyStatusSnapshot {
+	slot: KeySlot;
+	status: "healthy" | "disabled" | "cooling";
+	cooldownUntil?: number;
+}
+
+export interface ProviderStatusSnapshot {
+	providerId: string;
+	status: "healthy" | "cooling";
+	cooldownUntil?: number;
+	keys: KeyStatusSnapshot[];
+}
+
+export interface FailoverSnapshot {
+	providers: ProviderStatusSnapshot[];
+	active?: { providerId: string; model: string; keySlot: KeySlot };
+	visitedProviders: string[];
+	visitedKeys: Array<{ providerId: string; keySlot: KeySlot }>;
+}
+
+export interface ProviderPlan {
+	providerId: string;
+	model: string;
+}
+
+export interface ProviderFallbackRequest {
+	current: ProviderPlan;
+	unavailableProviderIds: readonly string[];
+}
+
+export interface FailoverProvider {
+	id: string;
+	backupKeyCount?: number;
+}
+
+export interface FailoverEngineOptions {
+	providers: readonly FailoverProvider[];
+	now: () => number;
+	nextProvider: (request: ProviderFallbackRequest) => ProviderPlan | undefined;
+}
+
+export type FailoverDecision =
+	| { kind: "switch-key" | "switch-model"; providerId: string; model: string; keySlot: KeySlot }
+	| { kind: "none" | "exhausted" };
+
+interface KeyState {
+	slot: KeySlot;
+	disabled: boolean;
+	cooldownUntil: number;
+}
+
+interface ProviderState {
+	id: string;
+	keys: KeyState[];
+	cooldownUntil: number;
+}
+
+interface ActiveAttempt {
+	plan: ProviderPlan;
+	keySlot: KeySlot;
+}
+
+const KEY_RATE_LIMIT_COOLDOWN_MS = 60_000;
+const PROVIDER_COOLDOWN_MS = 30_000;
+
+export function backupSlot(index: number): KeySlot {
+	return index === 0 ? "backup" : `backup-${index + 1}`;
+}
+
+export function backupIndexForSlot(slot: KeySlot): number | undefined {
+	if (slot === "primary") return undefined;
+	if (slot === "backup") return 0;
+	const ordinal = Number(slot.slice("backup-".length));
+	return Number.isInteger(ordinal) && ordinal >= 2 ? ordinal - 1 : undefined;
+}
+
+export class FailoverEngine {
+	private readonly providers: ProviderState[];
+	private readonly now: () => number;
+	private readonly nextProvider: FailoverEngineOptions["nextProvider"];
+	private visitedKeys = new Set<string>();
+	private visitedProviders = new Set<string>();
+	private active?: ActiveAttempt;
+	private decision: FailoverDecision = { kind: "none" };
+
+	constructor(options: FailoverEngineOptions) {
+		this.providers = options.providers.map((provider) => ({
+			id: provider.id,
+			keys: [
+				{ slot: "primary" as const, disabled: false, cooldownUntil: 0 },
+				...Array.from({ length: provider.backupKeyCount ?? 0 }, (_, index) => ({
+					slot: backupSlot(index),
+					disabled: false,
+					cooldownUntil: 0,
+				})),
+			],
+			cooldownUntil: 0,
+		}));
+		this.now = options.now;
+		this.nextProvider = options.nextProvider;
+	}
+
+	startTurn(initial: ProviderPlan): FailoverDecision {
+		this.visitedKeys = new Set();
+		this.visitedProviders = new Set();
+		this.active = undefined;
+
+		return this.selectPlan(initial, "switch-key") ?? this.selectFallback(initial);
+	}
+
+	observeFailure(observation: FailureObservation): FailoverDecision {
+		if (!this.active) return this.setDecision({ kind: "none" });
+
+		switch (classifyFailure(observation)) {
+			case "unauthorized":
+				{
+					const key = this.activeKey();
+					if (key) key.disabled = true;
+				}
+				return this.rotateKeyOrFallback();
+			case "rate-limited":
+				this.activeKey()!.cooldownUntil = this.now() + cooldownFor(observation, KEY_RATE_LIMIT_COOLDOWN_MS);
+				return this.rotateKeyOrFallback();
+			case "overloaded":
+			case "provider-error":
+			case "network":
+				this.activeProvider()!.cooldownUntil = this.now() + cooldownFor(observation, PROVIDER_COOLDOWN_MS);
+				return this.selectFallback(this.active.plan);
+			case "other":
+				return this.setDecision({ kind: "none" });
+		}
+	}
+
+	observeSuccess(): void {
+		if (!this.active) return;
+		const provider = this.activeProvider();
+		const key = this.activeKey();
+		if (!provider || !key) return;
+
+		provider.cooldownUntil = 0;
+		key.cooldownUntil = 0;
+	}
+
+	currentDecision(): FailoverDecision {
+		return this.decision;
+	}
+
+	resumeAttempt(attempt: FailoverAttempt): boolean {
+		const provider = this.provider(attempt.providerId);
+		const key = provider?.keys.find((candidate) => candidate.slot === attempt.keySlot);
+		if (!provider || !key) return false;
+		this.active = {
+			plan: { providerId: attempt.providerId, model: attempt.model },
+			keySlot: attempt.keySlot,
+		};
+		return true;
+	}
+
+	snapshot(): FailoverSnapshot {
+		const now = this.now();
+		return {
+			providers: this.providers.map((provider) => ({
+				providerId: provider.id,
+				...(provider.cooldownUntil > now
+					? { status: "cooling" as const, cooldownUntil: provider.cooldownUntil }
+					: { status: "healthy" as const }),
+				keys: provider.keys.map((key) => ({
+					slot: key.slot,
+					...(key.disabled
+						? { status: "disabled" as const }
+						: key.cooldownUntil > now
+							? { status: "cooling" as const, cooldownUntil: key.cooldownUntil }
+							: { status: "healthy" as const }),
+				})),
+			})),
+			...(this.active ? { active: { providerId: this.active.plan.providerId, model: this.active.plan.model, keySlot: this.active.keySlot } } : {}),
+			visitedProviders: [...this.visitedProviders],
+			visitedKeys: [...this.visitedKeys].map((value) => {
+				const separator = value.lastIndexOf(":");
+				return { providerId: value.slice(0, separator), keySlot: value.slice(separator + 1) as KeySlot };
+			}),
+		};
+	}
+
+	private rotateKeyOrFallback(): FailoverDecision {
+		const active = this.active;
+		if (!active) return this.setDecision({ kind: "none" });
+		const provider = this.activeProvider();
+		const key = provider && this.nextAvailableKey(provider);
+		if (!provider || !key) return this.selectFallback(active.plan);
+
+		return this.activate(active.plan, key, "switch-key");
+	}
+
+	private selectFallback(current: ProviderPlan): FailoverDecision {
+		const considered = new Set<string>();
+		for (let remaining = this.providers.length; remaining > 0; remaining -= 1) {
+			const plan = this.nextProvider({
+				current,
+				unavailableProviderIds: [...new Set([...this.visitedProviders, ...considered])],
+			});
+			if (!plan) break;
+			considered.add(plan.providerId);
+			const decision = this.selectPlan(plan, "switch-model");
+			if (decision) return decision;
+		}
+
+		return this.setDecision({ kind: "exhausted" });
+	}
+
+	private selectPlan(plan: ProviderPlan, kind: "switch-key" | "switch-model"): FailoverDecision | undefined {
+		const provider = this.provider(plan.providerId);
+		if (!provider || this.visitedProviders.has(provider.id) || provider.cooldownUntil > this.now()) return undefined;
+
+		const key = this.nextAvailableKey(provider);
+		if (!key) return undefined;
+
+		this.visitedProviders.add(provider.id);
+		return this.activate(plan, key, kind);
+	}
+
+	private activate(plan: ProviderPlan, key: KeyState, kind: "switch-key" | "switch-model"): FailoverDecision {
+		this.visitedKeys.add(this.keyId(plan.providerId, key.slot));
+		this.active = { plan, keySlot: key.slot };
+		return this.setDecision({ kind, providerId: plan.providerId, model: plan.model, keySlot: key.slot });
+	}
+
+	private nextAvailableKey(provider: ProviderState): KeyState | undefined {
+		const now = this.now();
+		return provider.keys.find(
+			(key) => !this.visitedKeys.has(this.keyId(provider.id, key.slot)) && !key.disabled && key.cooldownUntil <= now,
+		);
+	}
+
+	private activeProvider(): ProviderState | undefined {
+		return this.active && this.provider(this.active.plan.providerId);
+	}
+
+	private activeKey(): KeyState | undefined {
+		const provider = this.activeProvider();
+		return provider?.keys.find((key) => key.slot === this.active?.keySlot);
+	}
+
+	private provider(id: string): ProviderState | undefined {
+		return this.providers.find((provider) => provider.id === id);
+	}
+
+	private keyId(providerId: string, keySlot: KeySlot): string {
+		return `${providerId}:${keySlot}`;
+	}
+
+	private setDecision(decision: FailoverDecision): FailoverDecision {
+		this.decision = decision;
+		return decision;
+	}
+}
+
+function classifyFailure(observation: FailureObservation): FailureKind | "other" {
+	switch (observation.status) {
+		case 401:
+		case 403:
+			return "unauthorized";
+		case 429:
+			return "rate-limited";
+		case 529:
+			return "overloaded";
+		case 500:
+		case 502:
+		case 503:
+		case 504:
+			return "provider-error";
+	}
+
+	switch (observation.kind) {
+		case "overloaded":
+			return "overloaded";
+		case "provider-error":
+		case "network":
+			return observation.kind;
+		case "unauthorized":
+			return "unauthorized";
+		case "rate-limited":
+			return "rate-limited";
+		default:
+			return "other";
+	}
+}
+
+function cooldownFor(observation: FailureObservation, fallback: number): number {
+	return typeof observation.retryAfterMs === "number" && Number.isFinite(observation.retryAfterMs)
+		? Math.max(0, observation.retryAfterMs)
+		: fallback;
+}

package/src/failover-login.ts

@@ -0,0 +1,108 @@
+import type { AddBackupKeyRejection, AddBackupKeyResult } from "./add-backup-key.ts";
+import { isLiteralBackupKey, type AuthCatalog } from "./auth-catalog.ts";
+import type { NotificationLevel } from "./notification.ts";
+
+export interface FailoverLoginDependencies {
+	hasUI: boolean;
+	select: (title: string, options: string[]) => Promise<string | undefined>;
+	input: (title: string, placeholder?: string) => Promise<string | undefined>;
+	confirm: (title: string, message: string) => Promise<boolean>;
+	readCatalog: () => AuthCatalog;
+	writeBackupKey: (providerId: string, backupKey: string) => AddBackupKeyResult;
+	rebuild: () => void | Promise<void>;
+	notify: (message: string, level: NotificationLevel) => void;
+}
+
+/**
+ * Interactive `/failover-login` flow.
+ *
+ * An optional `[provider]` argument preselects the target; without one the user
+ * picks from the current `api_key` providers. The backup key always comes from an
+ * input dialog (never the command line), a confirmation gates the write, and the
+ * catalog is rebuilt on success so the new key is immediately usable.
+ * Notifications never include the key content.
+ */
+export async function runFailoverLogin(args: string, deps: FailoverLoginDependencies): Promise<void> {
+	if (!deps.hasUI) {
+		deps.notify("pi-failover: /failover-login requires an interactive session", "error");
+		return;
+	}
+
+	const requested = args.trim();
+	const catalog = deps.readCatalog();
+	const apiKeyProviders = catalog.providers.filter((provider) => provider.type === "api_key");
+
+	let providerId: string;
+	if (requested !== "") {
+		const match = catalog.providers.find((provider) => provider.provider === requested);
+		if (match === undefined) {
+			deps.notify(`pi-failover: unknown provider ${requested}`, "error");
+			return;
+		}
+		if (match.type !== "api_key") {
+			deps.notify(`pi-failover: ${requested} is an oauth provider; backup keys apply only to api_key credentials`, "error");
+			return;
+		}
+		providerId = requested;
+	} else {
+		if (apiKeyProviders.length === 0) {
+			deps.notify("pi-failover: no api_key providers configured", "error");
+			return;
+		}
+		const chosen = await deps.select(
+			"pi-failover: choose an api_key provider",
+			apiKeyProviders.map((provider) => provider.provider),
+		);
+		if (chosen === undefined) return;
+		providerId = chosen;
+	}
+
+	const backupKey = await deps.input("pi-failover: backup API key", "paste the backup key");
+	if (backupKey === undefined) return;
+	if (!isLiteralBackupKey(backupKey)) {
+		deps.notify("pi-failover: rejected backup key: must be non-empty and cannot start with ! or contain ${ or $VAR", "error");
+		return;
+	}
+
+	const confirmed = await deps.confirm(
+		`pi-failover: add a backup key to ${providerId}?`,
+		"The key will be written to auth.json and the failover catalog rebuilt.",
+	);
+	if (!confirmed) {
+		deps.notify("pi-failover: cancelled; auth.json unchanged", "info");
+		return;
+	}
+
+	const result = deps.writeBackupKey(providerId, backupKey);
+	if (!result.ok) {
+		deps.notify(`pi-failover: backup key not added for ${providerId}: ${describeRejection(result.reason)}`, "error");
+		return;
+	}
+
+	await deps.rebuild();
+	const count = result.backupKeys.length;
+	const noun = count === 1 ? "backup key" : "backup keys";
+	deps.notify(
+		result.changed
+			? `pi-failover: added a backup key to ${providerId} (${count} ${noun} configured)`
+			: `pi-failover: ${providerId} already has that backup key (${count} ${noun} configured)`,
+		"info",
+	);
+}
+
+function describeRejection(reason: AddBackupKeyRejection): string {
+	switch (reason) {
+		case "invalid-key":
+			return "invalid key";
+		case "unknown-provider":
+			return "provider not found in auth.json";
+		case "oauth-provider":
+			return "provider is oauth";
+		case "unreadable":
+			return "auth.json unreadable";
+		case "malformed":
+			return "auth.json malformed";
+		case "locked":
+			return "auth.json locked by another process";
+	}
+}