pi-keyrouter 0.2.3 → 0.3.1

package/index.ts

@@ -27,8 +27,9 @@
 //   # create ~/.pi/keyrouter.json with your provider keys
 //   /reload
 
-import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import type { ExtensionAPI, ExtensionUIContext } from "@earendil-works/pi-coding-agent";
 import { loadConfig, configPath } from "./config.ts";
+import { notifyRotation } from "./notification.ts";
 import {
 	initKeyStates,
 	isAvailable,
@@ -46,7 +47,7 @@ interface ProviderRuntime {
 export default function keyRouterExtension(pi: ExtensionAPI): void {
 	let config: KeyRouterConfig | undefined;
 	const runtimes = new Map<string, ProviderRuntime>();
-	let notify: ((text: string, level: "info" | "warning" | "error") => void) | undefined;
+	let uiCtx: ExtensionUIContext | undefined;
 	let activationNotified = false;
 
 	/**
@@ -75,7 +76,7 @@ export default function keyRouterExtension(pi: ExtensionAPI): void {
 	 */
 	async function activate(ctx: {
 		cwd: string;
-		ui: { notify: (t: string, l?: "info" | "warning" | "error") => void };
+		ui: ExtensionUIContext;
 		modelRegistry: { authStorage: { setRuntimeApiKey: (p: string, k: string) => void } };
 	}): Promise<void> {
 		// Load config once (reload clears it)
@@ -83,7 +84,7 @@ export default function keyRouterExtension(pi: ExtensionAPI): void {
 			config = loadConfig(ctx.cwd);
 		}
 		if (config.providers.length === 0) return;
-		notify = (text, level) => ctx.ui.notify(text, level);
+		uiCtx = ctx.ui;
 
 		const authStorage = ctx.modelRegistry.authStorage;
 		let newlyBootstrapped = 0;
@@ -160,8 +161,8 @@ export default function keyRouterExtension(pi: ExtensionAPI): void {
 		setKey(nextKey.value);
 		rt.currentIndex = nextIdx;
 
-		// Notify
-		if (notify && currentKey) {
+		// Notify with yellow box widget (falls back to plain notify)
+		if (currentKey && uiCtx) {
 			const event: RotationEvent = {
 				provider: providerName,
 				fromKey: currentKey.name,
@@ -170,11 +171,7 @@ export default function keyRouterExtension(pi: ExtensionAPI): void {
 				status,
 				attempt: rt.keys.reduce((a, k) => a + k.failures, 0),
 			};
-			notify(
-				`🔑 keyrouter: ${event.provider} — ${event.fromKey} → ${event.toKey} ` +
-					`(HTTP ${event.status}, ${event.reason})`,
-				"warning",
-			);
+			notifyRotation(uiCtx, event);
 		}
 		return true;
 	}
@@ -190,9 +187,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;
@@ -201,22 +216,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.
-			if (notify) {
+			// All keys exhausted — let pi surface the real error.
+			if (uiCtx) {
 				const failed = rt.keys.filter((k) => k.failures > 0).map((k) => k.name);
-				notify(
+				uiCtx.notify(
 					`🔑 keyrouter: ${providerName} — all keys exhausted (${failed.join(", ")}). ` +
-						`Letting pi surface the original HTTP ${event.status}.`,
+						`Letting pi surface the original error.`,
 					"error",
 				);
 			}
@@ -226,7 +237,7 @@ export default function keyRouterExtension(pi: ExtensionAPI): void {
 	pi.on("session_shutdown", () => {
 		runtimes.clear();
 		config = undefined;
-		notify = undefined;
+		uiCtx = undefined;
 		activationNotified = false;
 	});
 

package/notification.ts

@@ -0,0 +1,99 @@
+// =============================================================================
+// notification.ts — yellow Box widget for key rotation events
+// =============================================================================
+//
+// Uses the same pattern as pi-soly's notification.ts:
+//   ui.setWidget(key, (tui, theme) => Component, { placement })
+//
+// Why yellow? Rotations are not errors (the request will succeed on the next
+// key), so red toolErrorBg would be misleading. They're warnings, so we use
+// yellow (\x1b[43m) which isn't in pi's ThemeBg palette but is universally
+// supported across terminals (8-color ANSI, always available).
+//
+// Reset codes:
+//   \x1b[49m — reset background
+//   \x1b[39m — reset foreground (text color)
+
+import { Box, Spacer, Text } from "@earendil-works/pi-tui";
+import type { ExtensionUIContext } from "@earendil-works/pi-coding-agent";
+import type { RotationEvent } from "./types.ts";
+
+const WIDGET_KEY = "keyrouter-rotation";
+const AUTO_CLEAR_MS = 8000;
+
+// Standard ANSI yellow background (works in all terminals — 8-color minimum).
+// Truecolor would be nicer but detection is terminal-specific; 43m is safe.
+const YELLOW_BG = "\x1b[43m";
+const BLACK_FG = "\x1b[30m";
+const RESET_BG = "\x1b[49m";
+const RESET_FG = "\x1b[39m";
+const BOLD = "\x1b[1m";
+
+/** Wrap text in a yellow box (background + black text for contrast). */
+function yellowBg(text: string): string {
+	return `${YELLOW_BG}${BLACK_FG}${text}${RESET_BG}${RESET_FG}`;
+}
+
+/** Build the rotation widget. */
+function buildRotationBox(event: RotationEvent): Box {
+	const box = new Box(1, 0, (t) => yellowBg(t));
+
+	// Title: 🔑 keyrouter: provider — fromKey → toKey
+	const title = `${BOLD}🔑 keyrouter: ${event.provider} — ${event.fromKey} → ${event.toKey}${BOLD === "\x1b[1m" ? "\x1b[22m" : ""}`;
+	box.addChild(new Text(title, 1, 0));
+	box.addChild(new Spacer(1));
+
+	// Body: reason + status
+	const reasonText =
+		event.reason === "rate-limited"
+			? `Rate-limited (HTTP ${event.status}) — rotated to next key`
+			: `Unauthorized (HTTP ${event.status}) — skipping bad key`;
+	box.addChild(new Text(reasonText, 1, 0));
+
+	// Retry hint
+	box.addChild(new Text(`pi will retry with ${event.toKey}…`, 1, 0));
+
+	return box;
+}
+
+/**
+ * Show a yellow box widget for a key rotation event.
+ * Auto-clears after 8s.
+ */
+export function notifyRotation(ui: ExtensionUIContext, event: RotationEvent): void {
+	try {
+		ui.setWidget(
+			WIDGET_KEY,
+			() => buildRotationBox(event),
+			{ placement: "aboveEditor" },
+		);
+	} catch {
+		// setWidget may fail if UI not available (print mode) — fall back to notify
+		try {
+			ui.notify(
+				`🔑 keyrouter: ${event.provider} — ${event.fromKey} → ${event.toKey} (HTTP ${event.status}, ${event.reason})`,
+				"warning",
+			);
+		} catch {
+			// no UI at all — silent
+		}
+		return;
+	}
+	// Auto-clear after 8 seconds
+	setTimeout(() => {
+		try {
+			ui.setWidget(WIDGET_KEY, undefined);
+		} catch {
+			// session may have ended — ignore
+		}
+	}, AUTO_CLEAR_MS);
+}
+
+/** Clear the rotation widget manually. */
+export function clearRotationWidget(ui: ExtensionUIContext): void {
+	try {
+		ui.setWidget(WIDGET_KEY, undefined);
+	} catch {
+		// ignore
+	}
+}

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-keyrouter",
-  "version": "0.2.3",
+  "version": "0.3.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,6 +23,7 @@
     "index.ts",
     "config.ts",
     "rotation.ts",
+    "notification.ts",
     "types.ts"
   ],
   "keywords": [