pi-keyrouter 0.1.0 → 0.2.1

package/README.md

@@ -2,7 +2,7 @@
 
 **API key rotation for [pi-coding-agent](https://github.com/nicobailon/pi-coding-agent).**
 
-Multiple keys per provider · automatic 429/401 fallback · recursion-safe.
+Multiple keys per provider · automatic 429/401 fallback · native integration.
 
 ```bash
 pi install npm:pi-keyrouter
@@ -10,7 +10,7 @@ pi install npm:pi-keyrouter
 /reload
 ```
 
-When your model returns 429 (rate-limited) or 401 (unauthorized), the next key is picked automatically. pi sees a single successful response — retries are transparent.
+When your model returns 429 (rate-limited) or 401 (unauthorized), the next key is set via pi's native `authStorage.setRuntimeApiKey()`. pi's built-in retry then uses the new key automatically.
 
 ---
 
@@ -43,25 +43,40 @@ Add your provider config to `~/.pi/keyrouter.json`:
 
 ---
 
-## 🎯 How it works
+## 🎯 How it works (native integration)
 
-1. **Install** — extension loads, reads config, wraps `fetch`.
-2. **Request** — URL matches a provider → key picked, `Authorization: Bearer <key>` set.
-3. **On 429 / 401** — current key marked bad (cooldown `cooldownMs`), next key tried.
-4. **On 200** — response returned, key marked OK.
-5. **After `maxRetries`** — last failed response returned (so pi sees the real error).
+This extension uses pi's **native API key resolution** — no fetch hacks, no header manipulation.
 
-Recursion is bounded by `maxRetries` (default 3). No infinite loops.
+From the [pi SDK docs](https://github.com/nicobailon/pi-coding-agent):
+
+> API key resolution priority (handled by AuthStorage):
+> 1. **Runtime overrides (via `setRuntimeApiKey`, not persisted)** ← we use this
+> 2. Stored credentials in `auth.json`
+> 3. Environment variables
+> 4. Fallback resolver
+
+Flow:
+
+1. **session_start** — extension loads config, calls `authStorage.setRuntimeApiKey(provider, firstKey)` for each managed provider. Runtime override takes priority over auth.json.
+2. **Request** — pi makes the HTTP call with the runtime-overridden key.
+3. **after_provider_response (429/401/403)** — extension fires, calls `setRuntimeApiKey(provider, nextKey)`.
+4. **pi's built-in retry** — pi's retry logic (the "Retrying 3/3" you see in the UI) makes the next attempt, which now picks up the new runtime key.
+5. **Success or exhaustion** — if all keys fail, runtime override is cleared and pi surfaces the real error.
+
+### Why not fetch wrapping?
+
+An earlier version wrapped `globalThis.fetch`. It didn't work because the OpenAI SDK (used by pi-ai for z.ai and others) captures the `fetch` reference at client creation time, before extensions load. The SDK kept calling the original fetch, ignoring the wrapper.
+
+The native `setRuntimeApiKey` approach is cleaner: pi owns the HTTP layer, we only swap the key between attempts. No monkey-patching, no timing issues.
 
 ### What gets rotated
 
 | Status | Action |
 |---|---|
-| 200 | Return response, mark key OK |
-| 429 | Mark key `rate-limited` (cooldown), try next |
-| 401 / 403 | Mark key `unauthorized` (cooldown), try next |
-| 5xx / network | Don't mark key bad — try next, but no cooldown |
-| `maxRetries` exhausted | Return last failed response |
+| 200 | Key marked OK |
+| 429 | Current key marked `rate-limited` (cooldown), `setRuntimeApiKey(nextKey)` |
+| 401 / 403 | Current key marked `unauthorized` (cooldown), `setRuntimeApiKey(nextKey)` |
+| All keys exhausted | Runtime override cleared, pi surfaces real error |
 
 ---
 
@@ -142,7 +157,7 @@ API keys live in plain text in `keyrouter.json`. **Don't commit it.** Options:
 ## 🛠 Development
 
 ```bash
-bun test          # 34 tests
+bun test          # 33 tests
 bun run typecheck # tsc --noEmit
 ```
 
@@ -150,16 +165,15 @@ Monorepo layout:
 
 ```
 packages/pi-keyrouter/
-├── index.ts          — extension entry point
+├── index.ts          — extension entry point (native setRuntimeApiKey)
 ├── rotation.ts       — pure key-pick logic
-├── fetch-wrapper.ts  — fetch interceptor with retry
 ├── config.ts         — config loader
 ├── types.ts          — shared types
 └── tests/
-    ├── rotation.test.ts      — pure logic
-    ├── fetch-wrapper.test.ts — integration with mocked fetch
-    ├── config.test.ts        — config loader
-    └── smoke.test.ts         — load-time smoke test
+    ├── rotation.test.ts              — pure logic
+    ├── provider-resolution.test.ts   — provider name mapping
+    ├── config.test.ts                — config loader
+    └── smoke.test.ts                 — load-time smoke test
 ```
 
 ---

package/index.ts

@@ -1,122 +1,293 @@
 // =============================================================================
-// index.ts — pi-keyrouter extension entry point
+// index.ts — pi-keyrouter extension entry point (native setRuntimeApiKey)
 // =============================================================================
 //
+// HOW IT WORKS (native integration, no fetch hacks):
+//
+// 1. pi makes a request with the current API key
+// 2. Provider returns 429 (rate-limited) or 401/403 (unauthorized)
+// 3. `after_provider_response` event fires with the HTTP status
+// 4. We call ctx.modelRegistry.authStorage.setRuntimeApiKey(provider, nextKey)
+// 5. pi's BUILT-IN retry logic kicks in → next attempt uses the new key
+// 6. Repeat until a key succeeds or we exhaust our key pool
+//
+// This is the native integration point documented in the SDK:
+//   "API key resolution priority:
+//    1. Runtime overrides (via setRuntimeApiKey, not persisted)
+//    2. Stored credentials in auth.json
+//    3. Environment variables
+//    4. Fallback resolver"
+//
+// We only touch priority #1 (runtime override). auth.json is never modified.
+// On session end, runtime overrides vanish (not persisted) — clean slate
+// for next session, which is exactly what we want for 429 rate limits.
+//
 // Usage:
 //   pi install npm:pi-keyrouter
 //   # create ~/.pi/keyrouter.json with your provider keys
 //   /reload
-//
-// On load:
-// 1. Reads keyrouter config (project or user-level)
-// 2. Wraps globalThis.fetch with rotation logic
-// 3. On 429/401, retries with next key up to maxRetries
-// 4. Notifies user on every key switch via Box widget
-//
-// Provides /keyrouter command for status / disable / enable.
 
 import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
 import { loadConfig } from "./config.ts";
-import { installKeyRouter, type KeyRouterHandle } from "./fetch-wrapper.ts";
+import {
+	initKeyStates,
+	isAvailable,
+	markBad,
+	pickNextKey,
+} from "./rotation.ts";
+import type { KeyRouterConfig, RotationEvent, KeyState } from "./types.ts";
+
+interface ProviderRuntime {
+	keys: KeyState[];
+	/** Index of the key currently set via setRuntimeApiKey. -1 = none set yet. */
+	currentIndex: number;
+}
 
 export default function keyRouterExtension(pi: ExtensionAPI): void {
-	let handle: KeyRouterHandle | undefined;
-	let enabled = true;
-	let currentCwd = "";
-
-	function activate(cwd: string, notify: (text: string, level: string) => void): void {
-		if (handle) return;
-		const config = loadConfig(cwd);
-		if (config.providers.length === 0) {
-			return; // nothing to do
+	let config: KeyRouterConfig | undefined;
+	const runtimes = new Map<string, ProviderRuntime>();
+	let notify: ((text: string, level: "info" | "warning" | "error") => void) | undefined;
+	let activationNotified = false;
+
+	function ensureRuntime(providerName: string, cfg: KeyRouterConfig): ProviderRuntime | undefined {
+		let rt = runtimes.get(providerName);
+		if (rt) return rt;
+		const providerCfg = cfg.providers.find((p) => p.name === providerName);
+		if (!providerCfg) return undefined;
+		rt = {
+			keys: initKeyStates(providerCfg.keys),
+			currentIndex: -1,
+		};
+		runtimes.set(providerName, rt);
+		return rt;
+	}
+
+	/**
+	 * Activate the router: load config (once), bootstrap all providers.
+	 * Idempotent — safe to call on every before_agent_start. Only runs
+	 * the bootstrap the FIRST time for each provider.
+	 */
+	async function activate(ctx: {
+		cwd: string;
+		ui: { notify: (t: string, l?: "info" | "warning" | "error") => void };
+		modelRegistry: { authStorage: { setRuntimeApiKey: (p: string, k: string) => void } };
+	}): Promise<void> {
+		// Load config once (reload clears it)
+		if (!config) {
+			config = loadConfig(ctx.cwd);
+		}
+		if (config.providers.length === 0) return;
+		notify = (text, level) => ctx.ui.notify(text, level);
+
+		const authStorage = ctx.modelRegistry.authStorage;
+		let newlyBootstrapped = 0;
+		for (const p of config.providers) {
+			const providerName = resolveProviderName(p.name);
+			// Skip providers we've already bootstrapped
+			if (runtimes.has(providerName)) continue;
+			if (bootstrap(providerName)) {
+				const rt = runtimes.get(providerName);
+				if (rt && rt.currentIndex >= 0) {
+					const key = rt.keys[rt.currentIndex];
+					if (key) {
+						authStorage.setRuntimeApiKey(providerName, key.value);
+						newlyBootstrapped++;
+					}
+				}
+			}
 		}
-		handle = installKeyRouter(config, (event) => {
+		// Only notify on first activation (when we bootstrapped at least one)
+		if (newlyBootstrapped > 0 && !activationNotified) {
+			activationNotified = true;
+			ctx.ui.notify(
+				`🔑 keyrouter: active (${config.providers.length} provider(s), ${config.providers.reduce((a, p) => a + p.keys.length, 0)} keys)`,
+				"info",
+			);
+		}
+	}
+
+	/** Set the initial key for a provider on first use. */
+	function bootstrap(providerName: string): boolean {
+		const cfg = config;
+		if (!cfg) return false;
+		const rt = ensureRuntime(providerName, cfg);
+		if (!rt) return false;
+		if (rt.currentIndex >= 0) return true; // already bootstrapped
+		const idx = pickNextKey(rt.keys, 0, Date.now());
+		if (idx < 0) return false;
+		const key = rt.keys[idx];
+		if (!key) return false;
+		// We can't call setRuntimeApiKey here (no ctx), but we mark the index
+		// so the first after_provider_response knows where we are.
+		rt.currentIndex = idx;
+		return true;
+	}
+
+	/** Rotate to the next available key. Returns true if rotated. */
+	function rotate(
+		providerName: string,
+		reason: "rate-limited" | "unauthorized",
+		status: number,
+		setKey: (key: string) => void,
+	): boolean {
+		const cfg = config;
+		if (!cfg) return false;
+		const rt = runtimes.get(providerName);
+		if (!rt) return false;
+
+		// Mark current key as bad
+		const currentKey = rt.currentIndex >= 0 ? rt.keys[rt.currentIndex] : undefined;
+		if (currentKey) {
+			markBad(currentKey, reason, cfg.cooldownMs, Date.now());
+		}
+
+		// Find next available key (different from current)
+		const nextIdx = pickNextKey(rt.keys, rt.currentIndex + 1, Date.now());
+		if (nextIdx < 0 || nextIdx === rt.currentIndex) {
+			// No other key available
+			return false;
+		}
+		const nextKey = rt.keys[nextIdx];
+		if (!nextKey) return false;
+
+		// Set the new runtime key — pi's retry will use it
+		setKey(nextKey.value);
+		rt.currentIndex = nextIdx;
+
+		// Notify
+		if (notify && currentKey) {
+			const event: RotationEvent = {
+				provider: providerName,
+				fromKey: currentKey.name,
+				toKey: nextKey.name,
+				reason,
+				status,
+				attempt: rt.keys.reduce((a, k) => a + k.failures, 0),
+			};
 			notify(
 				`🔑 keyrouter: ${event.provider} — ${event.fromKey} → ${event.toKey} ` +
-					`(HTTP ${event.status}, attempt ${event.attempt})`,
+					`(HTTP ${event.status}, ${event.reason})`,
 				"warning",
 			);
-		});
-	}
-
-	function deactivate(): void {
-		if (handle) {
-			handle.disable();
-			handle = undefined;
 		}
+		return true;
 	}
 
 	pi.on("session_start", async (_event, ctx) => {
-		currentCwd = ctx.cwd;
-		activate(ctx.cwd, (text, level) =>
-			(ctx.ui.notify as (t: string, l?: string) => void)(text, level),
-		);
-		if (handle) {
-			ctx.ui.notify(
-				`🔑 keyrouter: active (${loadConfig(ctx.cwd).providers.length} provider(s))`,
-				"info",
-			);
+		await activate(ctx);
+	});
+
+	// Lazy bootstrap: also fire on every turn. This handles /reload (which
+	// does NOT re-fire session_start) and config changes mid-session.
+	// activate() is idempotent — only bootstraps once per provider.
+	pi.on("before_agent_start", async (_event, ctx) => {
+		await activate(ctx);
+	});
+
+	pi.on("after_provider_response", async (event, ctx) => {
+		if (!config) return;
+		if (event.status !== 429 && event.status !== 401 && event.status !== 403) return;
+
+		// Determine provider from current model
+		const model = ctx.model;
+		if (!model) return;
+		const providerName = resolveProviderName(model.provider);
+		const rt = runtimes.get(providerName);
+		if (!rt) return; // not a managed provider
+
+		const reason: "rate-limited" | "unauthorized" =
+			event.status === 429 ? "rate-limited" : "unauthorized";
+
+		const authStorage = ctx.modelRegistry.authStorage;
+		const rotated = rotate(providerName, reason, event.status, (key) => {
+			authStorage.setRuntimeApiKey(providerName, key);
+		});
+
+		if (!rotated) {
+			// All keys exhausted — clear runtime override so pi falls back to auth.json
+			// (which has the user's original key). pi will surface the real error.
+			if (notify) {
+				const failed = rt.keys.filter((k) => k.failures > 0).map((k) => k.name);
+				notify(
+					`🔑 keyrouter: ${providerName} — all keys exhausted (${failed.join(", ")}). ` +
+						`Letting pi surface the original HTTP ${event.status}.`,
+					"error",
+				);
+			}
 		}
 	});
 
 	pi.on("session_shutdown", () => {
-		deactivate();
+		runtimes.clear();
+		config = undefined;
+		notify = undefined;
+		activationNotified = false;
 	});
 
 	pi.registerCommand("keyrouter", {
-		description: "manage key rotation (status, enable, disable, reload)",
+		description: "manage key rotation (status, reload)",
 		handler: async (args, ctx) => {
 			const sub = args.trim().split(/\s+/)[0] ?? "status";
 			if (sub === "status") {
-				if (!handle) {
-					ctx.ui.notify("🔑 keyrouter: not active", "info");
+				// On-demand activation in case session_start/before_agent_start
+				// haven't fired yet (e.g. user ran /keyrouter status right after
+				// /reload without sending a prompt).
+				if (!config || runtimes.size === 0) {
+					await activate(ctx);
+				}
+				if (!config || runtimes.size === 0) {
+					ctx.ui.notify(
+						"🔑 keyrouter: not active — no providers in ~/.pi/keyrouter.json " +
+							"(checked cwd/.soly, cwd/.pi, cwd, ~/.soly, ~/.pi, ~)",
+						"warning",
+					);
 					return;
 				}
-				const snap = handle.getSnapshot();
-				const lines: string[] = [`🔑 keyrouter: ${enabled ? "active" : "disabled"}`];
-				for (const p of snap) {
-					lines.push(``);
-					lines.push(`  ${p.provider} (current: ${p.current})`);
-					for (const k of p.keys) {
-						const cooldown = k.cooldownRemainingMs > 0
-							? ` ⏱ ${Math.ceil(k.cooldownRemainingMs / 1000)}s`
-							: "";
+				const lines: string[] = [`🔑 keyrouter: active`];
+				for (const [providerName, rt] of runtimes) {
+					const current = rt.currentIndex >= 0 ? rt.keys[rt.currentIndex] : undefined;
+					lines.push("");
+					lines.push(`  ${providerName} (current: ${current?.name ?? "(none)"})`);
+					for (const k of rt.keys) {
+						const marker = k === current ? "→" : "•";
+						const avail = isAvailable(k, Date.now()) ? "" : " (cooldown)";
 						lines.push(
-							`    • ${k.name}  uses=${k.uses} fails=${k.failures} status=${k.lastStatus}${cooldown}`,
+							`    ${marker} ${k.name}  uses=0 fails=${k.failures} status=${k.lastStatus}${avail}`,
 						);
 					}
 				}
 				ctx.ui.notify(lines.join("\n"), "info");
 				return;
 			}
-			if (sub === "enable") {
-				if (!handle) {
-					activate(currentCwd, (text, level) =>
-						(ctx.ui.notify as (t: string, l?: string) => void)(text, level),
-					);
-				}
-				enabled = true;
-				ctx.ui.notify("🔑 keyrouter: enabled", "info");
-				return;
-			}
-			if (sub === "disable") {
-				deactivate();
-				enabled = false;
-				ctx.ui.notify("🔑 keyrouter: disabled (fetch restored)", "info");
-				return;
-			}
 			if (sub === "reload") {
-				deactivate();
-				activate(currentCwd, (text, level) =>
-					(ctx.ui.notify as (t: string, l?: string) => void)(text, level),
+				config = loadConfig(ctx.cwd);
+				runtimes.clear();
+				activationNotified = false;
+				ctx.ui.notify(
+					`🔑 keyrouter: reloaded (${config.providers.length} provider(s))`,
+					"info",
 				);
-				ctx.ui.notify("🔑 keyrouter: reloaded", "info");
 				return;
 			}
-			ctx.ui.notify(
-				"Usage: /keyrouter [status|enable|disable|reload]",
-				"info",
-			);
+			ctx.ui.notify("Usage: /keyrouter [status|reload]", "info");
 		},
 	});
+}
+
+/**
+ * Resolve the internal provider name that authStorage uses.
+ * The keyrouter config uses display names like "z-ai" but authStorage
+ * uses the canonical provider id like "zai". We try a few mappings.
+ */
+function resolveProviderName(displayName: string): string {
+	const lower = displayName.toLowerCase();
+	// Common mappings
+	const map: Record<string, string> = {
+		"z-ai": "zai",
+		"z.ai": "zai",
+		"open-router": "openrouter",
+		"openai": "openai",
+		"anthropic": "anthropic",
+	};
+	return map[lower] ?? displayName;
 }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-keyrouter",
-  "version": "0.1.0",
+  "version": "0.2.1",
   "description": "API key rotation for pi-coding-agent. Multiple keys per provider, automatic 429/401 fallback, max-retries guard.",
   "type": "module",
   "main": "index.ts",
@@ -23,7 +23,6 @@
     "index.ts",
     "config.ts",
     "rotation.ts",
-    "fetch-wrapper.ts",
     "types.ts"
   ],
   "keywords": [

package/fetch-wrapper.ts

@@ -1,228 +0,0 @@
-// =============================================================================
-// fetch-wrapper.ts — wraps global fetch with key-rotation logic
-// =============================================================================
-//
-// Replaces `globalThis.fetch` with a function that:
-// 1. Intercepts requests whose URL matches a configured provider.
-// 2. Picks the best available key (rotation logic in rotation.ts).
-// 3. Sets the Authorization header.
-// 4. On 429/401, marks the key as bad and retries with the next key
-//    (up to maxRetries).
-// 5. Calls onRotate on every key switch.
-//
-// On failure, the original response is returned (not a synthetic one) so
-// pi sees the real error if all retries fail.
-
-import {
-	initKeyStates,
-	isAvailable,
-	markBad,
-	markOk,
-	matchProvider,
-	pickNextKey,
-	recordUse,
-	waitForNextKey,
-} from "./rotation.ts";
-import type {
-	KeyRouterConfig,
-	KeyState,
-	ProviderConfig,
-	RotationEvent,
-} from "./types.ts";
-
-/** State tracked per provider. */
-interface ProviderState {
-	config: ProviderConfig;
-	keys: KeyState[];
-	preferredIndex: number;
-}
-
-export interface KeyRouterHandle {
-	/** Restore the original fetch and stop intercepting. */
-	disable: () => void;
-	/** Snapshot of current state (for /keyrouter status command). */
-	getSnapshot: () => KeyRouterSnapshot[];
-}
-
-export interface KeyRouterSnapshot {
-	provider: string;
-	current: string;
-	keys: Array<{
-		name: string;
-		uses: number;
-		failures: number;
-		lastStatus: string;
-		cooldownRemainingMs: number;
-	}>;
-}
-
-/**
- * Install the fetch wrapper. Returns a handle for disable / inspection.
- *
- * @param config — key router config
- * @param onRotate — called on every key switch (for UI notification)
- */
-export function installKeyRouter(
-	config: KeyRouterConfig,
-	onRotate: (event: RotationEvent) => void,
-): KeyRouterHandle {
-	// Capture original fetch BEFORE wrapping
-	const originalFetch = globalThis.fetch.bind(globalThis);
-
-	// Build per-provider state
-	const providerStates = new Map<string, ProviderState>();
-	for (const p of config.providers) {
-		providerStates.set(p.name, {
-			config: p,
-			keys: initKeyStates(p.keys),
-			preferredIndex: 0,
-		});
-	}
-
-	async function wrappedFetch(
-		input: string | URL | Request,
-		init?: RequestInit,
-	): Promise<Response> {
-		const url =
-			typeof input === "string"
-				? input
-				: input instanceof URL
-					? input.toString()
-					: input.url;
-		const matched = matchProvider(
-			Array.from(providerStates.values()).map((s) => s.config),
-			url,
-		);
-		if (!matched) {
-			return originalFetch(input, init);
-		}
-		const state = providerStates.get(matched.name);
-		if (!state || state.keys.length === 0) {
-			return originalFetch(input, init);
-		}
-
-		const now = Date.now();
-		const maxRetries = Math.max(1, config.maxRetries);
-		let attempt = 0;
-		let lastResponse: Response | undefined;
-		let lastError: unknown;
-		let lastPreferred = state.preferredIndex;
-
-		while (attempt < maxRetries) {
-			attempt += 1;
-			const idx = pickNextKey(state.keys, lastPreferred, now);
-			if (idx < 0) break;
-			const key = state.keys[idx];
-			if (!key) break;
-
-			// If the chosen key is on cooldown, wait briefly
-			if (!isAvailable(key, now)) {
-				const wait = waitForNextKey(state.keys, now);
-				if (wait > 0 && wait < 2000) {
-					await new Promise((r) => setTimeout(r, wait));
-				}
-			}
-
-			recordUse(key);
-			const initCopy = { ...(init ?? {}) };
-			const headers = new Headers(initCopy.headers ?? {});
-			headers.set("Authorization", `Bearer ${key.value}`);
-			initCopy.headers = headers;
-
-			let response: Response;
-			try {
-				response = await originalFetch(input, initCopy);
-			} catch (e) {
-				lastError = e;
-				// Network errors don't consume a retry budget
-				// (we'll loop back and try again)
-				lastPreferred = (idx + 1) % state.keys.length;
-				continue;
-			}
-
-			if (response.status === 429) {
-				markBad(key, "rate-limited", config.cooldownMs, Date.now());
-				lastResponse = response;
-				if (attempt >= maxRetries) break;
-				const nextIdx = pickNextKey(state.keys, idx + 1, Date.now());
-				if (nextIdx === idx) break;
-				const nextKey = state.keys[nextIdx];
-				if (nextKey) {
-					onRotate({
-						provider: matched.name,
-						fromKey: key.name,
-						toKey: nextKey.name,
-						reason: "rate-limited",
-						status: 429,
-						attempt,
-					});
-				}
-				state.preferredIndex = nextIdx;
-				lastPreferred = nextIdx;
-				continue;
-			}
-
-			if (response.status === 401 || response.status === 403) {
-				markBad(key, "unauthorized", config.cooldownMs, Date.now());
-				lastResponse = response;
-				if (attempt >= maxRetries) break;
-				const nextIdx = pickNextKey(state.keys, idx + 1, Date.now());
-				if (nextIdx === idx) break;
-				const nextKey = state.keys[nextIdx];
-				if (nextKey) {
-					onRotate({
-						provider: matched.name,
-						fromKey: key.name,
-						toKey: nextKey.name,
-						reason: "unauthorized",
-						status: response.status,
-						attempt,
-					});
-				}
-				state.preferredIndex = nextIdx;
-				lastPreferred = nextIdx;
-				continue;
-			}
-
-			// Success — mark ok and return
-			markOk(key);
-			state.preferredIndex = idx;
-			return response;
-		}
-
-		// All retries exhausted — return the last response if we have one
-		if (lastResponse) return lastResponse;
-		// Or re-throw the last network error
-		if (lastError !== undefined) throw lastError;
-		// Or fall through to original fetch
-		return originalFetch(input, init);
-	}
-
-	// Install wrapper
-	(globalThis as { fetch: typeof fetch }).fetch = wrappedFetch as typeof fetch;
-
-	function getSnapshot(): KeyRouterSnapshot[] {
-		const now = Date.now();
-		return Array.from(providerStates.values()).map((s) => {
-			const current = s.keys[s.preferredIndex];
-			return {
-				provider: s.config.name,
-				current: current?.name ?? "(none)",
-				keys: s.keys.map((k) => ({
-					name: k.name,
-					uses: k.uses,
-					failures: k.failures,
-					lastStatus: k.lastStatus,
-					cooldownRemainingMs:
-						k.cooldownUntil > now ? k.cooldownUntil - now : 0,
-				})),
-			};
-		});
-	}
-
-	function disable(): void {
-		(globalThis as { fetch: typeof fetch }).fetch = originalFetch;
-	}
-
-	return { disable, getSnapshot };
-}