pi-keyrouter 0.2.2 → 0.3.0

package/index.ts

@@ -49,16 +49,22 @@ export default function keyRouterExtension(pi: ExtensionAPI): void {
 	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);
+	/**
+	 * Get-or-create the runtime for a provider. Keyed by the RESOLVED
+	 * name (the authStorage id, e.g. "zai"), but populated from the
+	 * provider config passed in directly (avoids name-mismatch bugs).
+	 */
+	function ensureRuntime(
+		resolvedName: string,
+		providerCfg: { keys: ReadonlyArray<{ name: string; value: string }> },
+	): ProviderRuntime {
+		let rt = runtimes.get(resolvedName);
 		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);
+		runtimes.set(resolvedName, rt);
 		return rt;
 	}
 
@@ -82,15 +88,15 @@ export default function keyRouterExtension(pi: ExtensionAPI): void {
 		const authStorage = ctx.modelRegistry.authStorage;
 		let newlyBootstrapped = 0;
 		for (const p of config.providers) {
-			const providerName = resolveProviderName(p.name);
+			const resolvedName = resolveProviderName(p.name);
 			// Skip providers we've already bootstrapped
-			if (runtimes.has(providerName)) continue;
-			if (bootstrap(providerName)) {
-				const rt = runtimes.get(providerName);
+			if (runtimes.has(resolvedName)) continue;
+			if (bootstrap(resolvedName, p)) {
+				const rt = runtimes.get(resolvedName);
 				if (rt && rt.currentIndex >= 0) {
 					const key = rt.keys[rt.currentIndex];
 					if (key) {
-						authStorage.setRuntimeApiKey(providerName, key.value);
+						authStorage.setRuntimeApiKey(resolvedName, key.value);
 						newlyBootstrapped++;
 					}
 				}
@@ -107,11 +113,11 @@ export default function keyRouterExtension(pi: ExtensionAPI): void {
 	}
 
 	/** 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;
+	function bootstrap(
+		resolvedName: string,
+		providerCfg: { keys: ReadonlyArray<{ name: string; value: string }> },
+	): boolean {
+		const rt = ensureRuntime(resolvedName, providerCfg);
 		if (rt.currentIndex >= 0) return true; // already bootstrapped
 		const idx = pickNextKey(rt.keys, 0, Date.now());
 		if (idx < 0) return false;
@@ -184,9 +190,27 @@ export default function keyRouterExtension(pi: ExtensionAPI): void {
 		await activate(ctx);
 	});
 
-	pi.on("after_provider_response", async (event, ctx) => {
+	pi.on("message_end", async (event, ctx) => {
 		if (!config) return;
-		if (event.status !== 429 && event.status !== 401 && event.status !== 403) return;
+		const msg = event.message;
+		// Only intercept assistant error messages
+		if (msg.role !== "assistant" || msg.stopReason !== "error") return;
+		const errMsg = msg.errorMessage ?? "";
+		if (!errMsg) return;
+
+		// Detect error type from the message string.
+		// pi's error messages look like: "429 Usage limit reached..."
+		// or "401 Unauthorized" / "403 Forbidden".
+		let reason: "rate-limited" | "unauthorized" | null = null;
+		let status = 0;
+		if (/\b429\b|rate.?limit|too many requests/i.test(errMsg)) {
+			reason = "rate-limited";
+			status = 429;
+		} else if (/\b40[13]\b|unauthorized|forbidden/i.test(errMsg)) {
+			reason = "unauthorized";
+			status = errMsg.includes("401") ? 401 : 403;
+		}
+		if (!reason) return; // not a rotatable error
 
 		// Determine provider from current model
 		const model = ctx.model;
@@ -195,22 +219,18 @@ export default function keyRouterExtension(pi: ExtensionAPI): void {
 		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) => {
+		const rotated = rotate(providerName, reason, 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.
+			// All keys exhausted — let pi 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}.`,
+						`Letting pi surface the original error.`,
 					"error",
 				);
 			}

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-keyrouter",
-  "version": "0.2.2",
+  "version": "0.3.0",
   "description": "API key rotation for pi-coding-agent. Multiple keys per provider, automatic 429/401 fallback, max-retries guard.",
   "type": "module",
   "main": "index.ts",