pi-multi-account 1.2.0 → 1.4.0

package/CHANGELOG.md

@@ -5,6 +5,46 @@ All notable changes to this project are documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [1.4.0] - 2026-06-10
+
+### Fixed
+
+- **A single 401 no longer drops an account that still has valid tokens.** A 401 on
+  an OAuth account usually just means the access token needs a refresh (Pi refreshes
+  on the next call). Previously the first 401 permanently invalidated the account
+  (≈1-year cooldown until re-login) and yanked you onto another — often broken —
+  account. Now a refreshable account is given a brief cooldown and retried; it is
+  only marked dead after 3 consecutive 401s with no success in between. A
+  non-refreshable (API-key) 401 is still treated as immediately fatal.
+- Any successful response clears that account's 401 streak.
+
+### Added
+
+- Tests for transient-401 tolerance, the consecutive-401 kill threshold, and
+  success-resets-streak (suite now 17 tests).
+
+## [1.3.0] - 2026-06-10
+
+### Fixed
+
+- **Manual model/account selection is now respected.** Picking a model (e.g. Opus
+  on another account) no longer gets auto-yanked onto a different provider on the
+  next rate limit — the failover stays put and tells you, until you switch with
+  `/model` or `/multi-account next`. The pin auto-releases after a successful
+  response on that provider.
+- **No more self-resurrecting work.** All background resume timers were removed:
+  continuation now happens only synchronously inside an active turn, so Esc and
+  quitting always stop it. When every account is rate-limited the failover STOPS
+  and asks you to retry, instead of churning between exhausted accounts.
+- **No more "Agent is already processing" / "Cannot continue from message role:
+  assistant".** Continuations are sent only when the agent is idle and not aborting.
+
+### Added
+
+- Test suite (`npm test`) covering the failover edge cases: limit/401 failover,
+  all-accounts-exhausted stop, Esc/abort, manual-selection pinning, idle gating,
+  Anthropic OAuth shaping idempotency, and session shutdown. Wired into CI.
+
 ## [1.2.0] - 2026-06-10
 
 ### Added
@@ -76,6 +116,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Plaintext-free credential handling (SHA-256 fingerprints only); `0600`
   config/state files.
 
+[1.3.0]: https://github.com/Sarrius/pi-multi-account/releases/tag/v1.3.0
 [1.2.0]: https://github.com/Sarrius/pi-multi-account/releases/tag/v1.2.0
 [1.1.0]: https://github.com/Sarrius/pi-multi-account/releases/tag/v1.1.0
 [1.0.0]: https://github.com/Sarrius/pi-multi-account/releases/tag/v1.0.0

package/index.ts

@@ -113,6 +113,16 @@ const DEFAULT_INVALID_COOLDOWN_MS = 365 * 24 * 60 * 60 * 1000; // effectively "u
 // until the machine swaps itself to death.
 const ANTI_PINGPONG_MS = 60 * 1000; // don't switch straight back to the account we just left
 const MIN_AUTOCONTINUE_INTERVAL_MS = 15 * 1000; // floor between auto-continuations (CPU/network guard)
+// A 401 on an OAuth account usually just means "access token expired, refresh it" — Pi refreshes
+// on the next call. So a single 401 must NOT permanently kill a refreshable account (that was the
+// "it dropped me off an account that still had tokens" bug). Only kill it after this many 401s in a
+// row with no successful response in between. Non-refreshable (API-key) 401 is fatal immediately.
+// Bumped on every release. Printed at startup and in `/multi-account status` so you can verify
+// which version Pi actually loaded (a running Pi keeps the version it started with — /login and
+// /reload do NOT reload extension code; only a full restart does).
+const VERSION = "1.4.0";
+const MAX_CONSECUTIVE_AUTH_FAILURES = 3;
+const TRANSIENT_AUTH_COOLDOWN_MS = 60 * 1000; // brief skip after a 401 so the next call can refresh
 
 const ANTHROPIC_BASE = "anthropic";
 const CODEX_BASE = "openai-codex";
@@ -800,6 +810,9 @@ export default function piMultiAccount(pi: ExtensionAPI) {
 
 	const exhaustedUntilByProvider = new Map<string, number>(Object.entries(persistedState.exhaustedUntilByProvider ?? {}));
 	const invalidatedByProvider = new Map<string, InvalidationRecord>(Object.entries(persistedState.invalidatedByProvider ?? {}));
+	// Consecutive 401s per provider (in-memory, reset on any success). Used so a transient 401 on a
+	// refreshable OAuth account doesn't permanently kill it.
+	const consecutiveAuthFailures = new Map<string, number>();
 
 	// Discovered, authed, deduped provider ids in rotation order.
 	let rotation: string[] = [];
@@ -824,6 +837,11 @@ export default function piMultiAccount(pi: ExtensionAPI) {
 	let autoContinueTimer: ReturnType<typeof setTimeout> | undefined; // pending spaced continuation
 	let lastLeftProvider: string | undefined; // account we just failed away from (anti-ping-pong)
 	let lastLeftAt = 0;
+	// When the USER manually picks a model/account, we respect it: auto-failover will not yank
+	// them off it (that's the "I selected opus and it flipped to chatgpt" bug). selfModelSwitch
+	// marks our OWN setModel calls so the model_select event isn't mistaken for a manual pick.
+	let userSelectedProvider: string | undefined;
+	let selfModelSwitch = false;
 	// The thinking level the user intended for this turn. pi.setModel() re-clamps and
 	// persists the thinking level on every model switch, so without this it drifts
 	// downward across failovers ("thinking level keeps dropping"). We capture it before
@@ -897,6 +915,43 @@ export default function piMultiAccount(pi: ExtensionAPI) {
 		return invalidatedByProvider.has(provider);
 	}
 
+	/** True if this account can self-heal a 401 by refreshing its OAuth token. */
+	function isRefreshable(provider: string): boolean {
+		const entry = readAuthFile()[provider];
+		return !!entry && typeof entry.refresh === "string" && entry.refresh.length > 0;
+	}
+
+	/** A successful response → this account's auth is fine; clear its 401 streak. */
+	function noteAuthSuccess(provider: string) {
+		if (consecutiveAuthFailures.delete(provider)) {
+			/* had a streak, now cleared */
+		}
+	}
+
+	/**
+	 * Handle a 401/auth failure WITHOUT nuking an account that just needs a token refresh.
+	 * Refreshable (OAuth) account: first failures → short cooldown only, so the next attempt can
+	 * refresh; only after MAX_CONSECUTIVE_AUTH_FAILURES in a row do we mark it dead-until-relogin.
+	 * Non-refreshable (API key): a 401 is genuinely fatal, mark invalid at once.
+	 * Returns true if the account was permanently invalidated.
+	 */
+	function markAuthFailure(provider: string, reason: string): boolean {
+		if (!isRefreshable(provider)) {
+			markInvalid(provider, reason);
+			return true;
+		}
+		const n = (consecutiveAuthFailures.get(provider) ?? 0) + 1;
+		consecutiveAuthFailures.set(provider, n);
+		if (n >= MAX_CONSECUTIVE_AUTH_FAILURES) {
+			consecutiveAuthFailures.delete(provider);
+			markInvalid(provider, `${reason} (after ${n} consecutive 401s)`);
+			return true;
+		}
+		// Transient: brief cooldown so selection skips it for a moment; Pi refreshes on next use.
+		markExhausted(provider, TRANSIENT_AUTH_COOLDOWN_MS);
+		return false;
+	}
+
 	// ----- cooldowns --------------------------------------------------------
 
 	function pruneCooldowns() {
@@ -1099,11 +1154,22 @@ export default function piMultiAccount(pi: ExtensionAPI) {
 		return pool.sort((a, b) => a.remaining - b.remaining || a.rotIndex - b.rotIndex).map((s) => s.model);
 	}
 
-	async function switchToFallback(ctx: any, reason: string, cooldownMs = config.cooldownMs) {
+	async function switchToFallback(ctx: any, reason: string, cooldownMs = config.cooldownMs, manual = false) {
 		if (!config.enabled) return false;
 		const currentModel = ctx.model;
 		if (!currentModel) return false;
 
+		// Respect a manual model choice: if the user just picked this provider, do NOT auto-yank
+		// them onto another one — show the error and let them decide. Manual /multi-account next
+		// bypasses this (manual=true).
+		if (!manual && userSelectedProvider && currentModel.provider === userSelectedProvider) {
+			ctx.ui.notify(
+				`Provider failover: you selected ${currentModel.provider}/${currentModel.id} manually — staying on it (${reason.slice(0, 90)}). Use /model or /multi-account next to switch.`,
+				"warning",
+			);
+			return false;
+		}
+
 		markExhausted(currentModel.provider, cooldownMs);
 		lastLeftProvider = currentModel.provider;
 		lastLeftAt = Date.now();
@@ -1127,7 +1193,9 @@ export default function piMultiAccount(pi: ExtensionAPI) {
 		const from = ref(currentModel.provider, currentModel.id);
 		for (const fallback of candidates) {
 			const to = ref(fallback.provider, fallback.id);
+			selfModelSwitch = true; // our own switch — not a manual user pick
 			const ok = await pi.setModel(fallback);
+			selfModelSwitch = false;
 			if (!ok) {
 				// setModel failed → the account has no usable auth right now.
 				ctx.ui.notify(`Provider failover: ${to} has no usable auth, dropping from rotation`, "warning");
@@ -1152,35 +1220,31 @@ export default function piMultiAccount(pi: ExtensionAPI) {
 		return config.continuationPrompt.replaceAll("{from}", record.from).replaceAll("{to}", record.to).replaceAll("{reason}", record.reason);
 	}
 
-	/** Mark that the next agent run is our own failover continuation, then send it. */
-	function dispatchSelfContinuation(ctx: any, prompt: string) {
+	/**
+	 * Send our failover continuation — but ONLY when it is genuinely safe:
+	 *   - the user has not aborted (Esc), and the current op isn't aborting, and
+	 *   - the agent is idle (sending mid-turn throws "Agent is already processing" /
+	 *     "Cannot continue from message role: assistant").
+	 * Returns whether it actually sent. No background timer is ever used, so a turn is
+	 * always active for Esc to cancel — Esc/quit therefore always stop the chain.
+	 */
+	function dispatchSelfContinuation(ctx: any, prompt: string): boolean {
+		if (userAbortedChain || ctx.signal?.aborted || !ctx.isIdle()) return false;
 		lastAutoContinueAt = Date.now();
 		lastSentContinuationPrompt = prompt;
 		expectingSelfContinuation = true;
-		pi.sendUserMessage(prompt, ctx.isIdle() ? undefined : { deliverAs: "followUp" });
+		pi.sendUserMessage(prompt);
+		return true;
 	}
 
 	/**
-	 * Send an auto-continuation, but never faster than MIN_AUTOCONTINUE_INTERVAL_MS.
-	 * The spacing keeps a fully rate-limited rotation from pegging CPU/network and gives
-	 * the user a real window in which Esc actually sticks.
+	 * Continue after a successful failover switch — SYNCHRONOUSLY only.
+	 * Deliberately NOT a setTimeout: a deferred timer fires sendUserMessage when there is no
+	 * active turn for Esc to cancel, which is exactly how the chain escaped the user's control
+	 * and resurrected work on its own. Returns whether it sent.
 	 */
-	function scheduleAutoContinue(ctx: any, prompt: string) {
-		if (autoContinueTimer) {
-			clearTimeout(autoContinueTimer);
-			autoContinueTimer = undefined;
-		}
-		const wait = Math.max(0, MIN_AUTOCONTINUE_INTERVAL_MS - (Date.now() - lastAutoContinueAt));
-		if (wait === 0) {
-			dispatchSelfContinuation(ctx, prompt);
-			return;
-		}
-		autoContinueTimer = setTimeout(() => {
-			autoContinueTimer = undefined;
-			if (userAbortedChain || ctx.signal?.aborted) return; // user took over while we waited
-			dispatchSelfContinuation(ctx, prompt);
-		}, wait);
-		ctx.ui.notify(`Provider failover: next auto-continue in ~${Math.ceil(wait / 1000)}s (press Esc to cancel).`, "info");
+	function scheduleAutoContinue(ctx: any, prompt: string): boolean {
+		return dispatchSelfContinuation(ctx, prompt);
 	}
 
 	function clearPendingContinuation() {
@@ -1192,96 +1256,21 @@ export default function piMultiAccount(pi: ExtensionAPI) {
 		persist();
 	}
 
-	function nextPendingWakeDelayMs() {
-		if (!persistedState.pendingContinuationPrompt) return undefined;
-		const now = Date.now();
-		const lastProbe = lastProbeMap();
-		let bestWakeAt = Number.POSITIVE_INFINITY;
-		for (const provider of configuredProviders()) {
-			if (isInvalidated(provider)) continue;
-			const exhaustedUntil = exhaustedUntilByProvider.get(provider) ?? 0;
-			if (exhaustedUntil <= now) return 1000;
-			const probeDueAt = (lastProbe[provider] ?? 0) + config.probeCooldownMs;
-			bestWakeAt = Math.min(bestWakeAt, exhaustedUntil, probeDueAt);
-		}
-		if (!Number.isFinite(bestWakeAt)) return config.probeCooldownMs;
-		return Math.max(1000, Math.min(bestWakeAt - now, 2_147_483_647));
-	}
-
-	function schedulePendingWake(ctx?: any) {
-		if (ctx) latestCtx = ctx;
-		if (pendingWakeTimer) clearTimeout(pendingWakeTimer);
-		const delayMs = nextPendingWakeDelayMs();
-		if (delayMs === undefined) return;
-		pendingWakeTimer = setTimeout(() => {
-			pendingWakeTimer = undefined;
-			void attemptPendingResume();
-		}, delayMs);
-	}
-
 	function setPendingContinuation(ctx: any, reason: string) {
-		// Don't re-arm or re-notify if a pending resume is already queued — switchToFallback
-		// and agent_end can both reach here for the same exhaustion, and the wake timer is
-		// already running.
-		const alreadyPending = !!persistedState.pendingContinuationPrompt;
-		const current = ctx.model ? ref(ctx.model.provider, ctx.model.id) : ("unknown/model" as ModelRef);
-		const record: SwitchRecord = { from: current, to: "next-available/account" as ModelRef, reason, at: Date.now() };
-		persistedState = {
-			...persistedState,
-			pendingContinuationPrompt: persistedState.pendingContinuationPrompt || continuationPrompt(record),
-			pendingSince: persistedState.pendingSince || Date.now(),
-			pendingReason: reason,
-		};
+		// Every available account is rate-limited or unavailable right now. We deliberately do
+		// NOT arm a background timer to auto-resume later: such a timer fires sendUserMessage with
+		// no active turn for Esc to cancel and resurrects work on its own. Instead we STOP cleanly
+		// and tell the user — they retry by sending a message when an account has recovered.
+		const alreadyStopped = persistedState.pendingReason === reason;
+		persistedState = { ...persistedState, pendingReason: reason, pendingContinuationPrompt: undefined, pendingSince: undefined };
 		persist();
-		schedulePendingWake(ctx);
-		if (alreadyPending) return;
-		const delayMs = nextPendingWakeDelayMs();
+		if (alreadyStopped) return;
 		ctx.ui.notify(
-			`Provider failover: all accounts appear exhausted. Will automatically probe/resume in ~${Math.ceil((delayMs ?? config.probeCooldownMs) / 1000)}s if this Pi session stays open.`,
+			`Provider failover: every account is rate-limited or unavailable right now — stopped here. Send a message to retry once one recovers (check /multi-account status).`,
 			"warning",
 		);
 	}
 
-	async function attemptPendingResume() {
-		const ctx = latestCtx;
-		const prompt = persistedState.pendingContinuationPrompt;
-		if (!ctx || !prompt || !config.enabled || !config.autoContinue) return;
-		if (userAbortedChain) {
-			clearPendingContinuation(); // user took over — abandon the background resurrection
-			return;
-		}
-		if (autoContinuesThisPrompt >= config.maxAutoContinuesPerPrompt) {
-			clearPendingContinuation(); // task-level cap reached — stop resurrecting
-			return;
-		}
-		refreshDiscovery();
-		pruneCooldowns();
-		const candidates = findFallbackModels(ctx, ctx.model);
-		if (candidates.length === 0) {
-			schedulePendingWake(ctx);
-			return;
-		}
-		for (const candidate of candidates) {
-			const to = ref(candidate.provider, candidate.id);
-			const ok = await pi.setModel(candidate);
-			if (!ok) {
-				markInvalid(candidate.provider, "setModel failed on resume");
-				continue;
-			}
-			restoreDesiredThinking(); // keep the user's thinking level across the switch
-			setLastProbe(candidate.provider);
-			clearPendingContinuation();
-			// A genuine recovery after a real wait earns a fresh continuation budget so the
-			// agent can keep going whenever an account recovers; rapid flapping (resume that
-			// immediately re-limits) does NOT reset, so the cap still bounds a tight loop.
-			if (Date.now() - lastAutoContinueAt >= config.probeCooldownMs) autoContinuesThisPrompt = 0;
-			ctx.ui.notify(`Provider failover: resuming pending work on ${to}`, "warning");
-			dispatchSelfContinuation(ctx, prompt);
-			return;
-		}
-		schedulePendingWake(ctx);
-	}
-
 	// ----- error classification --------------------------------------------
 
 	function isAuthError(text: string) {
@@ -1328,6 +1317,9 @@ export default function piMultiAccount(pi: ExtensionAPI) {
 			exhaustedUntilByProvider.clear();
 			currentPromptSwitch = undefined;
 			autoContinuesThisPrompt = 0;
+			userAbortedChain = false;
+			userSelectedProvider = undefined;
+			consecutiveAuthFailures.clear();
 			if (pendingWakeTimer) {
 				clearTimeout(pendingWakeTimer);
 				pendingWakeTimer = undefined;
@@ -1340,7 +1332,8 @@ export default function piMultiAccount(pi: ExtensionAPI) {
 			return;
 		}
 		if (command === "next") {
-			await switchToFallback(ctx, "manual /multi-account next", 5 * 60 * 1000);
+			userSelectedProvider = undefined; // explicit request to move — drop any manual pin
+			await switchToFallback(ctx, "manual /multi-account next", 5 * 60 * 1000, true);
 			return;
 		}
 		if (command === "enable" || command === "disable") {
@@ -1357,7 +1350,7 @@ export default function piMultiAccount(pi: ExtensionAPI) {
 		const invalids = [...invalidatedByProvider.entries()].map(([p, r]) => `${p} (${r.reason.slice(0, 40)})`);
 		ctx.ui.notify(
 			[
-				`pi-multi-account: ${config.enabled ? "enabled" : "disabled"}${config.autoDiscover ? " · auto-discover ON" : " · auto-discover OFF"}`,
+				`pi-multi-account v${VERSION}: ${config.enabled ? "enabled" : "disabled"}${config.autoDiscover ? " · auto-discover ON" : " · auto-discover OFF"}`,
 				`Current: ${current}`,
 				`Rotation (${rotation.length}): ${rotation.join(" → ") || "none — log in to an account"}`,
 				`Registered login slots: ${[...registeredSlots].join(", ") || "(base accounts only)"}`,
@@ -1402,7 +1395,7 @@ export default function piMultiAccount(pi: ExtensionAPI) {
 		expectingSelfContinuation = false;
 		lastSentContinuationPrompt = "";
 		ctx.ui.notify(
-			`pi-multi-account loaded (${config.enabled ? "enabled" : "disabled"}). ${rotation.length} account(s) in rotation. Config: ${CONFIG_PATH}`,
+			`pi-multi-account v${VERSION} loaded (${config.enabled ? "enabled" : "disabled"}). ${rotation.length} account(s) in rotation. Config: ${CONFIG_PATH}`,
 			"info",
 		);
 	});
@@ -1455,14 +1448,36 @@ export default function piMultiAccount(pi: ExtensionAPI) {
 		refreshDiscovery(); // cheap: only re-scans when auth.json changed (new /login)
 	});
 
+	// Detect a MANUAL model/account selection by the user (vs our own failover setModel) and
+	// pin it, so auto-failover won't immediately yank them off it.
+	pi.on("model_select", (event) => {
+		if (selfModelSwitch) return; // our own failover switch — not a manual pick
+		const model = (event as any).model;
+		if (model?.provider) userSelectedProvider = model.provider;
+	});
+
 	pi.on("after_provider_response", async (event, ctx) => {
 		latestCtx = ctx;
 		if (!config.enabled) return;
 		if (userAbortedChain || ctx.signal?.aborted) return; // user is cancelling — don't fail over
 		const status = (event as any).status;
+		// A successful response proves this account's auth is fine → clear its 401 streak, and
+		// release a manual pin so normal auto-failover resumes for the user's chosen model.
+		if (status < 400 && ctx.model) {
+			noteAuthSuccess(ctx.model.provider);
+			if (ctx.model.provider === userSelectedProvider) userSelectedProvider = undefined;
+		}
 		if (status === 401) {
-			// Authorization is dead → drop this account, then move on.
-			if (ctx.model) markInvalid(ctx.model.provider, `HTTP 401`);
+			// A 401 on an OAuth account usually just needs a token refresh — do NOT kill it on the
+			// first one. markAuthFailure only invalidates after repeated 401s; otherwise it's a
+			// brief cooldown so Pi can refresh and we retry the SAME account, not abandon it.
+			if (ctx.model) {
+				const killed = markAuthFailure(ctx.model.provider, "HTTP 401");
+				if (!killed) {
+					ctx.ui.notify(`Provider failover: ${ctx.model.provider} got a 401 — will refresh and retry (not dropping it).`, "info");
+					return; // give the same account a chance to refresh on the next call
+				}
+			}
 			await switchToFallback(ctx, "HTTP 401 (auth invalid)");
 			return;
 		}
@@ -1480,7 +1495,13 @@ export default function piMultiAccount(pi: ExtensionAPI) {
 		lastErrorText = errorText;
 		if (currentPromptSwitch) return;
 		if (isAuthError(errorText)) {
-			if (ctx.model) markInvalid(ctx.model.provider, errorText.slice(0, 60));
+			if (ctx.model) {
+				const killed = markAuthFailure(ctx.model.provider, errorText.slice(0, 60));
+				if (!killed) {
+					ctx.ui.notify(`Provider failover: ${ctx.model.provider} hit a transient auth error — will refresh and retry (not dropping it).`, "info");
+					return; // refreshable account: let it refresh and retry rather than abandoning it
+				}
+			}
 			await switchToFallback(ctx, `auth invalid: ${errorText.slice(0, 100)}`);
 			return;
 		}
@@ -1510,7 +1531,7 @@ export default function piMultiAccount(pi: ExtensionAPI) {
 		if (userAbortedChain) return;
 
 		const errorText = lastErrorText || getAssistantErrorText((event as any).messages ?? []);
-		if (isAuthError(errorText) && ctx.model) markInvalid(ctx.model.provider, errorText.slice(0, 60));
+		if (isAuthError(errorText) && ctx.model) markAuthFailure(ctx.model.provider, errorText.slice(0, 60));
 		if (!isLimitError(errorText) && !isAuthError(errorText)) return;
 
 		// Task-level cap. Because this counter is no longer reset by our own re-prompts,
@@ -1536,9 +1557,10 @@ export default function piMultiAccount(pi: ExtensionAPI) {
 		}
 
 		if (currentPromptSwitch) {
-			autoContinuesThisPrompt++;
 			const prompt = continuationPrompt(currentPromptSwitch);
-			scheduleAutoContinue(ctx, prompt); // spaced + Esc-cancellable, not a tight loop
+			// Continue synchronously and only if it actually sent (agent idle, not aborted).
+			// Count the attempt only when we really sent, so the cap reflects real tries.
+			if (scheduleAutoContinue(ctx, prompt)) autoContinuesThisPrompt++;
 		}
 	});
 }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-multi-account",
-  "version": "1.2.0",
+  "version": "1.4.0",
   "description": "Automatic multi-account failover & rotation for Pi Agent across Anthropic (Claude), OpenAI/ChatGPT Codex, and Qwen/Alibaba. Auto-discovers authenticated accounts, grows the rotation on login, and drops accounts on logout, expiry, or quota/rate-limit errors.",
   "type": "module",
   "license": "MIT",
@@ -42,6 +42,7 @@
   ],
   "scripts": {
     "check": "tsc --noEmit",
+    "test": "node --test test/*.test.ts",
     "pack:check": "npm pack --dry-run",
     "prepublishOnly": "npm run check"
   },