pi-gemini-oauth-apikey 0.1.0

package/README.md

@@ -0,0 +1,87 @@
+# pi-gemini-oauth-apikey
+
+Multi-account **Google Cloud Code Assist (Gemini CLI) OAuth rotation pool** for
+[pi](https://pi.dev). Register N Google accounts as separate providers and
+auto-rotate on 429/`RESOURCE_EXHAUSTED` — get ~N × 1500 RPD of free Gemini quota.
+
+Complements [`pi-key-pool`](https://github.com/ssdiwu/pi-key-pool) (which rotates
+**API keys**): this package rotates **OAuth accounts** (Code Assist tier,
+~6× the per-account RPD, separate quota bucket).
+
+## What it does
+
+- Registers `gemini-pool-1` … `gemini-pool-N` as OAuth providers
+- Each is independently loggable via `/login`
+- On a 429/rate-limit error, marks the current account exhausted, switches to
+  the next available one, and **replays your last message automatically**
+- Time-based cooldown recovery (quota 60s, capacity 30s, network skipped)
+
+## Why
+
+`pi-multi-pass` supports multi-account OAuth rotation but its Gemini login
+(`loginGeminiCli`) was removed from upstream pi, so `/subs login google-gemini-cli-N`
+throws `is not a function`. This package reuses
+[`@qraxiss/pi-gemini-auth`](https://github.com/qraxiss/pi-gemini-auth) (which
+restored that OAuth flow as a standalone package) for the login/provider
+implementation, and adds the rotation layer on top. No pi-core patching.
+
+## Install
+
+```bash
+pi install npm:pi-gemini-oauth-apikey
+```
+
+Then restart pi. (You can now uninstall `@qraxiss/pi-gemini-auth` as an active
+extension — this package depends on it as a library instead.)
+
+## Configure
+
+Optional. Edit `~/.pi/agent/gemini-pool.json` (defaults to 4 slots if absent):
+
+```json
+{
+  "accounts": [
+    { "id": "1", "label": "Work Pro" },
+    { "id": "2", "label": "Personal Pro" },
+    { "id": "3", "label": "Side Pro" },
+    { "id": "4", "label": "Backup Pro" }
+  ],
+  "cooldown": { "quota": 60000, "capacity": 30000, "network": 0 }
+}
+```
+
+## Login each account
+
+```
+/login        → pick "Gemini Pool 1" → browser OAuth (account 1)
+/login        → pick "Gemini Pool 2" → account 2
+... (repeat per account)
+/model gemini-pool-1/gemini-2.5-flash
+```
+
+That's it. When `gemini-pool-1` hits a 429, the turn transparently retries on
+`gemini-pool-2`, and so on. When all are exhausted you get a warning and a
+cooldown timer.
+
+## How it works
+
+```
+before_agent_start  →  remember last user prompt
+agent_end           →  if stopReason=="error" && isRateLimit(msg):
+                         markExhausted(current); pickNext() → setModel + sendUserMessage
+```
+
+`failover.ts` is a pure, unit-tested state machine (round-robin + cooldown +
+auth-gated availability). Provider/OAuth/streaming come from qraxiss.
+
+## Tests
+
+```bash
+node --test --experimental-strip-types test/*.test.ts
+```
+
+Stdlib only (`node:test`) — no test dependencies.
+
+## License
+
+MIT

package/package.json

@@ -0,0 +1,19 @@
+{
+  "name": "pi-gemini-oauth-apikey",
+  "version": "0.1.0",
+  "type": "module",
+  "description": "Multi-account Google Gemini (Cloud Code Assist OAuth) rotation pool for pi — register N accounts, auto-rotate on 429. Depends on @qraxiss/pi-gemini-auth for the OAuth/provider implementation.",
+  "license": "MIT",
+  "author": "vforvaick",
+  "keywords": ["pi-package", "pi-extension", "gemini", "oauth", "rotation", "pool", "multi-account"],
+  "repository": { "type": "git", "url": "git+https://github.com/vforvaick/pi-gemini-oauth-apikey.git" },
+  "dependencies": {
+    "@qraxiss/pi-gemini-auth": "^0.5.1"
+  },
+  "scripts": {
+    "test": "node --test --experimental-strip-types test/*.test.ts"
+  },
+  "pi": {
+    "extensions": ["./src/index.ts"]
+  }
+}

package/src/config.ts

@@ -0,0 +1,57 @@
+// ponytail: minimal config loader. ~/.pi/agent/gemini-pool.json optional;
+// defaults to 4 slots so it works zero-config for a typical multi-Pro setup.
+import { readFileSync, existsSync } from "node:fs";
+import { homedir } from "node:os";
+import { join } from "node:path";
+
+export interface Account {
+	id: string; // provider id: gemini-pool-<id>
+	label: string; // /login display name
+}
+
+export interface PoolConfig {
+	accounts: Account[];
+	cooldown?: { capacity?: number; quota?: number; network?: number };
+}
+
+export function configPath(): string {
+	return join(homedir(), ".pi", "agent", "gemini-pool.json");
+}
+
+export function defaultConfig(): PoolConfig {
+	return {
+		accounts: [1, 2, 3, 4].map((i) => ({
+			id: `${i}`,
+			label: `Gemini Pool ${i}`,
+		})),
+	};
+}
+
+export function providerName(account: Account): string {
+	return `gemini-pool-${account.id}`;
+}
+
+export function normalize(raw: unknown): PoolConfig {
+	const r = (raw ?? {}) as { accounts?: unknown; cooldown?: unknown };
+	const accounts: Account[] = Array.isArray(r.accounts)
+		? r.accounts
+				.map((a, i) => {
+					const obj = (a ?? {}) as { id?: string; label?: string };
+					const id = String(obj.id ?? i + 1);
+					return { id, label: obj.label || `Gemini Pool ${id}` };
+				})
+				.filter((a) => a.id)
+		: [];
+	if (accounts.length === 0) return defaultConfig();
+	return { accounts, cooldown: (r.cooldown ?? {}) as PoolConfig["cooldown"] };
+}
+
+export function loadConfig(): PoolConfig {
+	const path = configPath();
+	if (!existsSync(path)) return defaultConfig();
+	try {
+		return normalize(JSON.parse(readFileSync(path, "utf-8")));
+	} catch {
+		return defaultConfig(); // malformed → safe default, don't crash pi
+	}
+}

package/src/failover.ts

@@ -0,0 +1,95 @@
+// ponytail: pure state machine, zero pi deps, fully unit-testable.
+// Cooldown + round-robin rotation over a fixed member list.
+
+export type CooldownReason = "capacity" | "quota" | "network";
+
+export interface CooldownConfig {
+	capacity: number; // overloaded / 503 / 529 — transient
+	quota: number; // 429 / RESOURCE_EXHAUSTED — standard recovery
+	network: number; // network errors — don't blame the key (0 = no cooldown)
+}
+
+export const DEFAULT_COOLDOWN: CooldownConfig = {
+	capacity: 30_000,
+	quota: 60_000,
+	network: 0,
+};
+
+const RATE_LIMIT_PATTERNS = [
+	/429/,
+	/resource_exhausted/i,
+	/quota/i,
+	/rate.?limit/i,
+	/too many requests/i,
+	/overloaded/i,
+	/capacity/i,
+	/usage.?limit/i,
+	/limit.*reached/i,
+];
+
+export function isRateLimitError(msg: string): boolean {
+	return RATE_LIMIT_PATTERNS.some((p) => p.test(msg));
+}
+
+export function classifyError(msg: string): CooldownReason {
+	if (/\b(503|529)\b|overloaded|high demand|capacity/i.test(msg)) return "capacity";
+	if (/network|timeout|econn|socket/i.test(msg)) return "network";
+	return "quota";
+}
+
+export interface FailoverPool {
+	members: string[];
+	cooldown: CooldownConfig;
+	exhausted: Map<string, { reason: CooldownReason; until: number }>;
+}
+
+export function createPool(
+	members: string[],
+	cooldown: CooldownConfig = DEFAULT_COOLDOWN,
+): FailoverPool {
+	return { members, cooldown, exhausted: new Map() };
+}
+
+export function markExhausted(
+	pool: FailoverPool,
+	member: string,
+	reason: CooldownReason,
+	now: number,
+): void {
+	const cd = pool.cooldown[reason];
+	if (cd <= 0) return; // network: don't blame the member
+	pool.exhausted.set(member, { reason, until: now + cd });
+}
+
+export function isAvailable(
+	pool: FailoverPool,
+	member: string,
+	now: number,
+	hasAuth: boolean,
+): boolean {
+	if (!hasAuth) return false;
+	const e = pool.exhausted.get(member);
+	if (!e) return true;
+	if (now >= e.until) {
+		pool.exhausted.delete(member); // cooldown elapsed → member recovers
+		return true;
+	}
+	return false;
+}
+
+// Round-robin from the member after `current`, first available wins. null if all down.
+export function pickNext(
+	pool: FailoverPool,
+	current: string | null,
+	now: number,
+	hasAuthFn: (m: string) => boolean,
+): string | null {
+	if (pool.members.length === 0) return null;
+	const startIdx = current ? pool.members.indexOf(current) : -1;
+	for (let i = 1; i <= pool.members.length; i++) {
+		const idx = (startIdx + i) % pool.members.length;
+		const m = pool.members[idx];
+		if (isAvailable(pool, m, now, hasAuthFn(m))) return m;
+	}
+	return null;
+}

package/src/index.ts

@@ -0,0 +1,96 @@
+// pi-gemini-oauth-apikey: register N Google Cloud Code Assist (Gemini CLI)
+// OAuth accounts as separate providers and auto-rotate on 429.
+//
+// Reuses @qraxiss/pi-gemini-auth for the OAuth/provider implementation (rule 5:
+// already-installed dependency). This package only adds multi-account rotation.
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import {
+	GEMINI_CLI_API,
+	streamSimpleGoogleGeminiCli,
+} from "@qraxiss/pi-gemini-auth/src/gemini-cli-provider";
+import { geminiOAuthOverride } from "@qraxiss/pi-gemini-auth/src/gemini-oauth";
+import { GEMINI_CLI_MODELS } from "@qraxiss/pi-gemini-auth/src/models";
+import {
+	createPool,
+	markExhausted,
+	pickNext,
+	isRateLimitError,
+	classifyError,
+	DEFAULT_COOLDOWN,
+	type CooldownConfig,
+} from "./failover.ts";
+import { loadConfig, providerName, type Account } from "./config.ts";
+
+const BASE_URL = "https://cloudcode-pa.googleapis.com";
+
+// Clone qraxiss models under a pool provider id (keep id/name/cost etc, swap provider).
+function cloneModelsFor(provider: string) {
+	return GEMINI_CLI_MODELS.map((m) => ({ ...m, provider }));
+}
+
+export default function geminiPool(pi: ExtensionAPI) {
+	const config = loadConfig();
+	if (config.accounts.length === 0) return;
+
+	const cooldown: CooldownConfig = {
+		capacity: config.cooldown?.capacity ?? DEFAULT_COOLDOWN.capacity,
+		quota: config.cooldown?.quota ?? DEFAULT_COOLDOWN.quota,
+		network: config.cooldown?.network ?? DEFAULT_COOLDOWN.network,
+	};
+
+	const members = config.accounts.map(providerName);
+	const pool = createPool(members, cooldown);
+	const memberSet = new Set(members);
+
+	// Register each account as its own OAuth provider. /login lists them by label.
+	for (const account of config.accounts) {
+		const name = providerName(account);
+		pi.registerProvider(name, {
+			name: account.label,
+			api: GEMINI_CLI_API,
+			baseUrl: BASE_URL,
+			oauth: { ...geminiOAuthOverride, name: account.label },
+			streamSimple: streamSimpleGoogleGeminiCli,
+			models: cloneModelsFor(name),
+		});
+	}
+
+	let lastUserPrompt: string | null = null;
+
+	pi.on("before_agent_start", async (event: any) => {
+		lastUserPrompt = event?.prompt ?? lastUserPrompt;
+	});
+
+	pi.on("agent_end", async (event: any, ctx: any) => {
+		const messages = event?.messages;
+		if (!Array.isArray(messages) || messages.length === 0) return;
+		const last = messages[messages.length - 1];
+		if (last?.role !== "assistant" || last.stopReason !== "error") return;
+		const errorMsg: string = last.errorMessage || "";
+		if (!errorMsg || !isRateLimitError(errorMsg)) return;
+
+		const currentProvider = ctx?.model?.provider;
+		if (!currentProvider || !memberSet.has(currentProvider)) return; // not our pool
+
+		markExhausted(pool, currentProvider, classifyError(errorMsg), Date.now());
+
+		const hasAuth = (m: string) => ctx.modelRegistry?.authStorage?.hasAuth(m) ?? false;
+		const next = pickNext(pool, currentProvider, Date.now(), hasAuth);
+		if (!next) {
+			ctx?.ui?.notify?.(
+				`[gemini-pool] All ${members.length} members rate-limited. Retry in ~${Math.round(
+					cooldown.quota / 1000,
+				)}s.`,
+				"warning",
+			);
+			return;
+		}
+
+		const nextModel = ctx?.modelRegistry?.find?.(next, ctx.model?.id);
+		if (!nextModel) return;
+		const switched = await pi.setModel?.(nextModel);
+		if (switched !== false && lastUserPrompt) {
+			pi.sendUserMessage?.(lastUserPrompt); // replay the failed turn
+		}
+	});
+}

package/test/config.test.ts

@@ -0,0 +1,33 @@
+import { test } from "node:test";
+import assert from "node:assert/strict";
+import { normalize, defaultConfig, providerName } from "../src/config.ts";
+
+test("defaultConfig: 4 accounts, sequential ids, distinct provider names", () => {
+	const d = defaultConfig();
+	assert.equal(d.accounts.length, 4);
+	assert.deepEqual(
+		d.accounts.map(providerName),
+		["gemini-pool-1", "gemini-pool-2", "gemini-pool-3", "gemini-pool-4"],
+	);
+});
+
+test("normalize: explicit accounts preserved with labels", () => {
+	const c = normalize({ accounts: [{ id: "work", label: "Work Pro" }, { id: "8" }] });
+	assert.equal(c.accounts.length, 2);
+	assert.equal(c.accounts[0].label, "Work Pro");
+	assert.equal(providerName(c.accounts[1]), "gemini-pool-8");
+	// missing label gets generated default
+	assert.match(c.accounts[1].label, /Gemini Pool 8/);
+});
+
+test("normalize: empty/garbage input falls back to default (never zero accounts)", () => {
+	assert.equal(normalize({ accounts: [] }).accounts.length, 4);
+	assert.equal(normalize({}).accounts.length, 4);
+	assert.equal(normalize(null).accounts.length, 4);
+	assert.equal(normalize("garbage").accounts.length, 4);
+});
+
+test("normalize: cooldown passed through when present", () => {
+	const c = normalize({ accounts: [{ id: "1" }], cooldown: { quota: 120000 } });
+	assert.equal(c.cooldown?.quota, 120000);
+});

package/test/failover.test.ts

@@ -0,0 +1,72 @@
+// ponytail: stdlib only (node:test) — no vitest dep for a ~120-line package.
+import { test } from "node:test";
+import assert from "node:assert/strict";
+import {
+	createPool, markExhausted, isAvailable, pickNext,
+	isRateLimitError, classifyError, DEFAULT_COOLDOWN,
+} from "../src/failover.ts";
+
+const NOW = 1_000_000;
+const has = () => true;
+const hasNot = (_: string) => false;
+
+test("isRateLimitError: 429 / quota / RESOURCE_EXHAUSTED → true; unrelated → false", () => {
+	assert.equal(isRateLimitError("429 Too Many Requests"), true);
+	assert.equal(isRateLimitError("Quota exceeded for metric x"), true);
+	assert.equal(isRateLimitError("RESOURCE_EXHAUSTED"), true);
+	assert.equal(isRateLimitError("connection refused"), false);
+	assert.equal(isRateLimitError(""), false);
+});
+
+test("classifyError: capacity vs quota vs network", () => {
+	assert.equal(classifyError("503 high demand"), "capacity");
+	assert.equal(classifyError("overloaded"), "capacity");
+	assert.equal(classifyError("ETIMEDOUT network reset"), "network");
+	assert.equal(classifyError("429 quota exceeded"), "quota");
+});
+
+test("fresh pool: authed member available, unauthed not", () => {
+	const p = createPool(["a", "b", "c"]);
+	assert.equal(isAvailable(p, "a", NOW, true), true); // member + authed
+	assert.equal(isAvailable(p, "a", NOW, false), false); // member but not authed → skip
+});
+
+test("markExhausted makes member unavailable until cooldown elapses", () => {
+	const p = createPool(["a", "b"]);
+	markExhausted(p, "a", "quota", NOW);
+	assert.equal(isAvailable(p, "a", NOW, true), false); // immediately after
+	assert.equal(isAvailable(p, "a", NOW + DEFAULT_COOLDOWN.quota - 1, true), false); // just before
+	assert.equal(isAvailable(p, "a", NOW + DEFAULT_COOLDOWN.quota, true), true); // at expiry → recovered
+	assert.equal(p.exhausted.has("a"), false, "expired entry should be cleared on probe");
+});
+
+test("pickNext rotates past current, skipping exhausted and unauthed", () => {
+	const p = createPool(["a", "b", "c"]);
+	markExhausted(p, "b", "quota", NOW);
+	// current=a → next available is c (b exhausted)
+	assert.equal(pickNext(p, "a", NOW, has), "c");
+	// current=c → wraps to a
+	assert.equal(pickNext(p, "c", NOW, has), "a");
+	// current=null → first available
+	assert.equal(pickNext(p, null, NOW, has), "a");
+});
+
+test("pickNext returns null when every member is down", () => {
+	const p = createPool(["a", "b"]);
+	markExhausted(p, "a", "quota", NOW);
+	markExhausted(p, "b", "quota", NOW);
+	assert.equal(pickNext(p, "a", NOW, has), null);
+});
+
+test("pickNext skips unauthenticated members", () => {
+	const p = createPool(["a", "b", "c"]);
+	const authedOnly = (m: string) => m === "b";
+	assert.equal(pickNext(p, "a", NOW, authedOnly), "b");
+});
+
+test("network reason (cooldown 0) never marks exhausted", () => {
+	const p = createPool(["a"]);
+	markExhausted(p, "a", "network", NOW);
+	assert.equal(p.exhausted.has("a"), false);
+	assert.equal(isAvailable(p, "a", NOW, true), true);
+});