diff --git a/.gitignore b/.gitignore index 10f2b8a..41caa3c 100644 --- a/.gitignore +++ b/.gitignore @@ -33,3 +33,5 @@ gemini-debug-*.log # Finder (MacOS) folder config .DS_Store +.atl/ +.atl diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..fcae114 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,19 @@ +# OpenCode Agent Instructions + +## Commands +- **Install dependencies:** `bun install` +- **Run tests:** `bun test` (runs all tests). To run a specific test: `bun test path/to/test.ts` +- **Build project:** `bun run build` (uses `tsup` to compile ESM). +- **Check node import compatibility:** `bun run smoke:node-import` + +## Architecture & Conventions +- **Opencode Plugin**: This is an Opencode authentication plugin for Gemini. +- **Entrypoints**: `index.ts` is the main export. `src/plugin.ts` sets up the loader and `fetch` interception. +- **Multi-Account**: The plugin supports multiple Gemini accounts via `AccountPool` (`src/plugin/account-pool.ts`). Accounts are selected using a rotation strategy (cooldown-aware LRU with quota fallback). +- **Storage**: Auth state is maintained in-memory and dynamically added accounts persist to `~/.config/opencode/gemini-auth.json` via `src/plugin/config-store.ts`. (Note: During testing, it isolates to `gemini-auth.test.json`). +- **Quota Handling**: `gquota` tool dynamically resolves Google Cloud project IDs and aggregates metrics across multiple accounts. + +## Workflow +- **Spec-Driven Development (SDD)**: This repo uses SDD. Look for specs and task artifacts in `openspec/` (or via Engram memory) before making architectural changes. +- **Strict TDD**: Write tests for any new features or bug fixes. `bun test` must pass before considering a task complete. Keep tests next to source files (e.g., `feature.test.ts`). +- **Imports**: Ensure all necessary functions are exported from `src/plugin/types.ts` if they need to be shared across modules. diff --git a/README.md b/README.md index 2e57a1c..9040521 100644 --- a/README.md +++ b/README.md @@ -90,6 +90,63 @@ or via environment variables: You can also set `OPENCODE_GEMINI_PROJECT_ID`, `GOOGLE_CLOUD_PROJECT`, or `GOOGLE_CLOUD_PROJECT_ID` to supply the project ID via environment variables. +### Multi-Account Support + +This plugin supports using multiple Google accounts simultaneously to distribute quota usage and avoid `429 Quota Exhausted` errors. + +When an account exhausts its quota for a requested model, the plugin places that account on a temporary cooldown and seamlessly retries your request with the next healthiest account in the background. + +#### Adding Accounts via UI (Recommended) + +1. Open your terminal and run `opencode auth login`. +2. Choose **Google**. +3. Select **Add Gemini Account**. +4. Complete the OAuth flow in your browser. + +Accounts added via the UI are automatically persisted across Opencode restarts to `~/.config/opencode/gemini-auth.json`. + +#### Managing Accounts + +You can view the status of all your configured accounts by selecting **Manage Gemini Accounts** in the `opencode auth login` menu. This allows you to check health scores, active cooldowns, and manually enable or disable specific accounts. + +#### Configuring Accounts Manually (Fallback) + +If you prefer to configure your accounts manually or deploy the plugin in a headless environment, you can define them directly in your `opencode.json` under `provider.google.options.accounts`: + +```json +{ + "provider": { + "google": { + "options": { + "accounts": [ + { + "id": "work", + "email": "user@company.com", + "refreshToken": "your-refresh-token-1", + "projectId": "prod-project" + }, + { + "id": "personal", + "email": "user@gmail.com", + "refreshToken": "your-refresh-token-2" + } + ] + } + } + } +} +``` + +*Note: Refresh tokens can be obtained from the `~/.config/opencode/gemini-auth.json` file after logging in at least once, or extracted from Opencode's credential store.* + +#### Rotation Strategy + +The plugin uses a hybrid **Cooldown-aware LRU with Quota Fallback** strategy: +1. **Cooldowns:** Ignores any account that recently received a `429` error for the requested model. +2. **Quota preference:** Prefers accounts with >20% remaining quota (based on `/gquota` data). +3. **Health Scoring:** Weighs accounts based on their historical success rate, average latency, and cooldown frequency. +4. **LRU Tiebreaker:** Uses the least recently used account to ensure fair round-robin distribution among healthy accounts. + ### Proxy If your network requires an HTTP proxy for Google API calls, set diff --git a/openspec/changes/multi-account-rotation/tasks.md b/openspec/changes/multi-account-rotation/tasks.md new file mode 100644 index 0000000..2d7faf8 --- /dev/null +++ b/openspec/changes/multi-account-rotation/tasks.md @@ -0,0 +1,52 @@ +# Tasks: multi-account-rotation + +Chain strategy: stacked-to-main | PR 1 of 5 + +## Phase 1: Foundation (PR 1, ~300 lines) + +- [x] T1: types.ts — Add GeminiAccount, AccountState, HealthScore, CooldownState, AccountPoolConfig, RotationStrategy +- [x] T2: health.ts — HealthTracker class, O(1) incremental scoring (success 40%, quota 30%, latency 15%, cooldown 15%) +- [x] T3: rotation.ts — selectBestAccount() filter→health-weight→quota→LRU +- [x] T4: health.test.ts — Score monotonicity tests +- [x] T5: rotation.test.ts — All-filter scenarios + +## Phase 2: Core Pool (PR 2, ~350 lines) + +- [ ] T6: account-pool.ts — AccountPool class (select, cooldown, reportSuccess/Failure, get/add/remove/toggle) +- [ ] T7: account-pool.ts — Per-account refresh lock via Map +- [ ] T8: account-pool.ts — Pool-of-one wrapper for backwards compat +- [ ] T9: cache.ts — Re-key auth cache by account ID +- [ ] T10: account-pool.test.ts — Selection, cooldown isolation, refresh lock, pool-of-one + +## Phase 3: Auth & Project Context (PR 3, ~300 lines) + +- [ ] T11: token.ts — refreshAccessTokenForAccount() +- [ ] T12: token.test.ts — Per-account refresh lock tests +- [ ] T13: project/utils.ts — buildProjectCacheKeyForAccount() +- [ ] T14: project/context.ts — Cache keyed by account ID, ensureProjectContextForAccount() +- [ ] T15: project/context.test.ts — Cache isolation per account +- [ ] T16: provider.ts — Per-account project ID resolution +- [ ] T17: oauth-authorize.ts — createOAuthAuthorizeMethodForAccount() + +## Phase 4: Plugin Integration (PR 4, ~300 lines) + +- [ ] T18: plugin.ts — Loader detects accounts[], builds AccountPool, fetch calls pool.select() +- [ ] T19: plugin.ts — fetch calls pool.refreshAccount(), ensureProjectContextForAccount() +- [ ] T20: plugin.test.ts — Multi-account init, fetch dispatch, backwards compat + +## Phase 5: Retry & Quota (PR 4 continued) + +- [ ] T21: retry/index.ts — Cooldown keys include account ID namespace +- [ ] T22: retry/index.ts — Terminal 429 sets per-account cooldown, delegates to pool.select() +- [ ] T23: retry/index.ts — All-exhausted error message lists accounts with reasons +- [ ] T24: quota.ts — Per-account quota sections +- [ ] T25: quota.ts — Aggregate summary header +- [ ] T26: quota.test.ts — Per-account formatting, aggregate, single-account compat + +## Phase 6: Account Management (PR 5, ~350 lines) + +- [ ] T27: account-manager.ts — Add Gemini Account auth method +- [ ] T28: account-manager.ts — Manage Gemini Accounts (list, toggle) +- [ ] T29: account-manager.ts — Remove account flow +- [ ] T30: oauth-authorize.ts — Extend OAuth for multi-account context +- [ ] T31: account-manager.test.ts — Add/remove/toggle, OAuth integration, config persistence diff --git a/openspec/config.yaml b/openspec/config.yaml new file mode 100644 index 0000000..755f3ca --- /dev/null +++ b/openspec/config.yaml @@ -0,0 +1,33 @@ +# openspec/config.yaml +schema: spec-driven + +context: | + Tech stack: TypeScript (ESNext, strict mode), Bun runtime, ESM modules + Architecture: OpenCode plugin for Gemini OAuth auth — PKCE flow, token refresh, request proxying to Gemini Code Assist API, quota management + Testing: Bun built-in test runner (bun:test), 11 test files, no coverage tool configured + Style: TypeScript with strict mode, verbatim module syntax, no external linter/formatter configured + Build: tsup bundles to ESM, target node20, inlines external dependencies + CI: GitHub Actions — release pipeline via bun install + tsup build + npm publish + +rules: + proposal: + - Include rollback plan for risky changes + specs: + - Use Given/When/Then for scenarios + - Use RFC 2119 keywords (MUST, SHALL, SHOULD, MAY) + design: + - Include sequence diagrams for complex flows + - Document architecture decisions with rationale + tasks: + - Group by phase, use hierarchical numbering + - Keep tasks completable in one session + apply: + - Follow existing code patterns + tdd: true + test_command: "bun test" + verify: + test_command: "bun test" + build_command: "bun run build" + coverage_threshold: 0 + archive: + - Warn before merging destructive deltas diff --git a/src/plugin.ts b/src/plugin.ts index af39f14..6d39aff 100644 --- a/src/plugin.ts +++ b/src/plugin.ts @@ -1,9 +1,12 @@ import { GEMINI_PROVIDER_ID } from "./constants"; import { geminiFetch } from "./fetch"; -import { createOAuthAuthorizeMethod } from "./plugin/oauth-authorize"; +import { AccountPool } from "./plugin/account-pool"; +import { AccountManager } from "./plugin/account-manager"; import { accessTokenExpired, isOAuthAuth } from "./plugin/auth"; +import { createOAuthAuthorizeMethod, createAddAccountAuthMethod } from "./plugin/oauth-authorize"; import { resolveCachedAuth } from "./plugin/cache"; -import { ensureProjectContext, retrieveUserQuota } from "./plugin/project"; +import { loadAccountsFromDisk, saveAccountsToDisk } from "./plugin/config-store"; +import { ensureProjectContext, ensureProjectContextForAccount, retrieveUserQuota } from "./plugin/project"; import { createGeminiQuotaTool, GEMINI_QUOTA_TOOL_NAME, @@ -11,9 +14,11 @@ import { import { isGeminiDebugEnabled, logGeminiDebugMessage, startGeminiDebugRequest } from "./plugin/debug"; import { maybeShowGeminiCapacityToast, maybeShowGeminiTestToast } from "./plugin/notify"; import { + resolveAccountsFromProvider, resolveConfiguredProjectId, resolveConfiguredProjectIdFromClient, resolveConfiguredProjectIdFromConfig, + resolveProjectIdForAccount, } from "./plugin/provider"; import { isGenerativeLanguageRequest, @@ -23,8 +28,9 @@ import { transformGeminiResponse, } from "./plugin/request"; import { fetchWithRetry } from "./plugin/retry"; -import { refreshAccessToken } from "./plugin/token"; +import { refreshAccessToken, refreshAccessTokenForAccount } from "./plugin/token"; import type { + GeminiAccount, GetAuth, LoaderResult, OAuthAuthDetails, @@ -43,6 +49,7 @@ Do not call other tools. let latestGeminiAuthResolver: GetAuth | undefined; let latestGeminiConfiguredProjectId: string | undefined; let latestGeminiUserAgentModel: string | undefined; +let latestGeminiPool: AccountPool | undefined; /** * Registers the Gemini OAuth provider for Opencode, handling auth, request rewriting, @@ -77,6 +84,7 @@ export const GeminiCLIOAuthPlugin = async ( getAuthResolver: () => latestGeminiAuthResolver, getConfiguredProjectId: () => latestGeminiConfiguredProjectId, getUserAgentModel: () => latestGeminiUserAgentModel, + getPool: () => latestGeminiPool, }), }, auth: { @@ -92,6 +100,42 @@ export const GeminiCLIOAuthPlugin = async ( normalizeProviderModelCosts(provider); const thinkingConfigDefaults = resolveThinkingConfigDefaults(provider); + // Initialize AccountPool from disk, then config accounts[], or pool-of-one + if (!latestGeminiPool) { + const diskAccounts = await loadAccountsFromDisk() || []; + const configAccounts = resolveAccountsFromProvider(provider) || []; + + // Merge accounts from both sources, deduplicating by ID (config accounts take precedence for overriding) + const mergedAccountsMap = new Map(); + + for (const acc of diskAccounts) { + mergedAccountsMap.set(acc.id, acc); + } + for (const acc of configAccounts) { + mergedAccountsMap.set(acc.id, acc); + } + + const accounts = Array.from(mergedAccountsMap.values()); + + if (accounts && accounts.length > 0) { + latestGeminiPool = new AccountPool({ accounts, strategy: "health-weighted" }); + for (const account of accounts) { + // Hydrate the matching account from the global getAuth() if it's the primary one + if (auth.refresh && account.refreshToken === auth.refresh) { + latestGeminiPool.updateAuth(account.id, auth); + } + } + } else { + const singleAccount: GeminiAccount = { + id: auth.refresh.split("|")[0] || "default", + refreshToken: auth.refresh, + enabled: true, + }; + latestGeminiPool = AccountPool.fromSingleAccount(singleAccount); + latestGeminiPool.updateAuth(singleAccount.id, auth); + } + } + return { apiKey: "", async fetch(input, init) { @@ -99,75 +143,176 @@ export const GeminiCLIOAuthPlugin = async ( return geminiFetch(input, init); } - const latestAuth = await getAuth(); - if (!isOAuthAuth(latestAuth)) { + // Select account from pool + const requestTarget = parseGenerativeLanguageRequest(input); + const model = requestTarget?.effectiveModel; + const url = toUrlString(input); + + let selected = latestGeminiPool?.select(model, url); + if (!selected) { return geminiFetch(input, init); } - let authRecord = resolveCachedAuth(latestAuth); - if (accessTokenExpired(authRecord)) { - const refreshed = await refreshAccessToken(authRecord, client); - if (!refreshed) { - return geminiFetch(input, init); + const maxAttempts = latestGeminiPool ? latestGeminiPool.count() : 1; + let attempt = 0; + let lastResponse: Response | null = null; + let lastTransformedModel: string | undefined = undefined; + let lastDebugContext: any = null; + let lastProjectContext: any = null; + + let lastTransformedStreaming: boolean = false; + + while (attempt < maxAttempts) { + attempt++; + + const currentAccountId = selected.account.id; + + // Refresh token if needed + if (accessTokenExpired(selected.auth)) { + try { + const refreshed = await refreshAccessTokenForAccount(latestGeminiPool!, currentAccountId, client); + if (refreshed) { + selected = latestGeminiPool!.getAccount(currentAccountId); + } + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + const isTimeout = message.includes("timeout"); + logGeminiDebugMessage(`Token refresh for account ${currentAccountId} failed: ${message}`); + + // Cooldown account and try next + latestGeminiPool!.cooldownAccount( + currentAccountId, + model ?? "all", + 60000, + isTimeout ? "REFRESH_TIMEOUT" : "REFRESH_ERROR" + ); + selected = latestGeminiPool!.select(model, url); + if (!selected) break; + continue; + } + + if (!selected?.auth.access) { + // If this account can't refresh, cooldown and try next + latestGeminiPool!.cooldownAccount(currentAccountId, model ?? "all", 60000, "AUTH_FAILED"); + selected = latestGeminiPool!.select(model, url); + if (!selected) break; + continue; + } } - authRecord = refreshed; - } - if (!authRecord.access) { - return geminiFetch(input, init); + const requestUserAgentModel = model; + if (requestUserAgentModel) { + latestGeminiUserAgentModel = requestUserAgentModel; + } + + // Resolve project context for this account + const configuredProjectId = resolveProjectIdForAccount( + selected.account, + await resolveLatestConfiguredProjectId(provider), + ); + const projectContext = await ensureProjectContextForAccount( + selected.auth, + client, + selected.account.id, + configuredProjectId, + requestUserAgentModel, + ); + lastProjectContext = projectContext; + + await maybeShowGeminiTestToast(client, projectContext.effectiveProjectId); + await maybeLogAvailableQuotaModels( + selected.auth.access!, + projectContext.effectiveProjectId, + requestUserAgentModel, + ); + + const transformed = prepareGeminiRequest( + input, + init, + selected.auth.access!, + projectContext.effectiveProjectId, + thinkingConfigDefaults, + ); + lastTransformedModel = transformed.requestedModel; + lastTransformedStreaming = transformed.streaming; + + const debugContext = startGeminiDebugRequest({ + originalUrl: toUrlString(input), + resolvedUrl: toUrlString(transformed.request), + method: transformed.init.method, + headers: transformed.init.headers, + body: transformed.init.body, + streaming: transformed.streaming, + projectId: projectContext.effectiveProjectId, + }); + lastDebugContext = debugContext; + + /** + * Retry transport/429 failures while preserving the requested model. + * We intentionally do not auto-downgrade model tiers to avoid misleading users. + */ + const response = await fetchWithRetry(transformed.request, transformed.init, selected.account.id); + lastResponse = response; + + // Report result to pool (handles success, failure, and 401 circuit breaker) + latestGeminiPool!.reportResult(selected.account.id, response.status); + + if (response.ok) { + break; // Success! Exit the loop. + } else { + // Check for terminal 429 to cooldown this account + if (response.status === 429 && response.headers.get("X-Gemini-Terminal-429") === "true") { + const reason = response.headers.get("X-Gemini-429-Reason") ?? "UNKNOWN"; + + // Default 60s cooldown (C-02) or honor Retry-After + let durationMs = 60000; + const retryAfter = response.headers.get("Retry-After"); + if (retryAfter) { + const parsed = parseInt(retryAfter, 10); + if (!isNaN(parsed)) { + durationMs = parsed * 1000; + } else { + const date = Date.parse(retryAfter); + if (!isNaN(date)) { + durationMs = Math.max(0, date - Date.now()); + } + } + } + + latestGeminiPool!.cooldownAccount(selected.account.id, model ?? "all", durationMs, reason); + + // Try to select another account + const nextAccount = latestGeminiPool!.select(model, url); + if (!nextAccount || nextAccount.account.id === selected.account.id) { + // No other healthy accounts available, break and return the 429 + break; + } + selected = nextAccount; + continue; // Loop again with the new account + } + + // If it's not a terminal 429 (e.g. 500 server error, 400 bad request), + // we don't rotate accounts. We just return the error. + break; + } } - const configuredProjectId = await resolveLatestConfiguredProjectId(provider); - const requestTarget = parseGenerativeLanguageRequest(input); - const requestUserAgentModel = requestTarget?.effectiveModel; - if (requestUserAgentModel) { - latestGeminiUserAgentModel = requestUserAgentModel; + // Fallback if loop finishes without a response + if (!lastResponse) { + return geminiFetch(input, init); } - const projectContext = await ensureProjectContextOrThrow( - authRecord, - client, - configuredProjectId, - requestUserAgentModel, - ); - await maybeShowGeminiTestToast(client, projectContext.effectiveProjectId); - await maybeLogAvailableQuotaModels( - authRecord.access, - projectContext.effectiveProjectId, - requestUserAgentModel, - ); - const transformed = prepareGeminiRequest( - input, - init, - authRecord.access, - projectContext.effectiveProjectId, - thinkingConfigDefaults, - ); - const debugContext = startGeminiDebugRequest({ - originalUrl: toUrlString(input), - resolvedUrl: toUrlString(transformed.request), - method: transformed.init.method, - headers: transformed.init.headers, - body: transformed.init.body, - streaming: transformed.streaming, - projectId: projectContext.effectiveProjectId, - }); - - /** - * Retry transport/429 failures while preserving the requested model. - * We intentionally do not auto-downgrade model tiers to avoid misleading users. - */ - const response = await fetchWithRetry(transformed.request, transformed.init); + await maybeShowGeminiCapacityToast( client, - response, - projectContext.effectiveProjectId, - transformed.requestedModel, + lastResponse, + lastProjectContext?.effectiveProjectId ?? "unknown", + lastTransformedModel, ); return transformGeminiResponse( - response, - transformed.streaming, - debugContext, - transformed.requestedModel, + lastResponse, + lastTransformedStreaming, + lastDebugContext, + lastTransformedModel, ); }, }; @@ -181,6 +326,18 @@ export const GeminiCLIOAuthPlugin = async ( getUserAgentModel: () => latestGeminiUserAgentModel, }), }, + createAddAccountAuthMethod(() => { + if (!latestGeminiPool) { + throw new Error("Gemini plugin not yet loaded (pool is undefined)"); + } + return new AccountManager({ + pool: latestGeminiPool, + client, + onConfigUpdate: (accounts) => { + saveAccountsToDisk(accounts).catch(console.error); + } + }); + }), { provider: GEMINI_PROVIDER_ID, label: "Manually enter API Key", @@ -325,3 +482,20 @@ async function maybeLogAvailableQuotaModels( `Code Assist models visible via quota buckets (${projectId}): ${modelIds.join(", ")}`, ); } + +/** + * Returns the current AccountPool instance (for testing and quota tool access). + */ +export function getLatestGeminiPool(): AccountPool | undefined { + return latestGeminiPool; +} + +/** + * Resets plugin module state for testing. + */ +export function resetPluginState(): void { + latestGeminiPool = undefined; + latestGeminiAuthResolver = undefined; + latestGeminiConfiguredProjectId = undefined; + latestGeminiUserAgentModel = undefined; +} diff --git a/src/plugin/account-manager.test.ts b/src/plugin/account-manager.test.ts new file mode 100644 index 0000000..fed074a --- /dev/null +++ b/src/plugin/account-manager.test.ts @@ -0,0 +1,120 @@ +import { describe, expect, test, mock } from "bun:test"; +import { AccountManager } from "./account-manager"; +import { AccountPool } from "./account-pool"; +import type { GeminiAccount, PluginClient } from "./types"; +import { resolveCachedAuthForAccount } from "./cache"; +import { invalidateProjectContextCacheForAccount } from "./project/context"; + +describe("AccountManager", () => { + const mockClient: PluginClient = { + auth: { set: mock(() => Promise.resolve()) } + }; + + test("addAccount adds to pool and notifies config update", async () => { + const pool = new AccountPool({ accounts: [] }); + let notifiedAccounts: GeminiAccount[] = []; + const manager = new AccountManager({ + pool, + client: mockClient, + onConfigUpdate: (accounts) => { notifiedAccounts = accounts; } + }); + + const account: GeminiAccount = { id: "test", refreshToken: "ref", enabled: true }; + await manager.addAccount(account); + + expect(pool.count()).toBe(1); + expect(notifiedAccounts.length).toBe(1); + expect(notifiedAccounts[0].id).toBe("test"); + }); + + test("removeAccount removes from pool and notifies", () => { + const account: GeminiAccount = { id: "test", refreshToken: "ref", enabled: true }; + const pool = new AccountPool({ accounts: [account] }); + let notifiedAccounts: GeminiAccount[] | null = null; + const manager = new AccountManager({ + pool, + client: mockClient, + onConfigUpdate: (accounts) => { notifiedAccounts = accounts; } + }); + + manager.removeAccount("test"); + + expect(pool.count()).toBe(0); + expect(notifiedAccounts?.length).toBe(0); + }); + + test("toggleAccount enables/disables account and notifies", () => { + const account: GeminiAccount = { id: "test", refreshToken: "ref", enabled: true }; + const pool = new AccountPool({ accounts: [account] }); + let notifiedAccounts: GeminiAccount[] | null = null; + const manager = new AccountManager({ + pool, + client: mockClient, + onConfigUpdate: (accounts) => { notifiedAccounts = accounts; } + }); + + manager.toggleAccount("test", false); + + expect(pool.getAccount("test")?.account.enabled).toBe(false); + expect(notifiedAccounts?.[0].enabled).toBe(false); + }); + + test("listAccounts returns status for all accounts", () => { + const pool = new AccountPool({ + accounts: [ + { id: "acc1", email: "1@test.com", refreshToken: "ref1", enabled: true }, + { id: "acc2", refreshToken: "ref2", enabled: false } + ] + }); + + // Simulate cooldown on acc1 + pool.cooldownAccount("acc1", "gemini-test", 10000, "TEST_COOLDOWN"); + + const manager = new AccountManager({ pool, client: mockClient }); + const list = manager.listAccounts(); + + expect(list.length).toBe(2); + + const acc1 = list.find(a => a.id === "acc1"); + expect(acc1?.email).toBe("1@test.com"); + expect(acc1?.enabled).toBe(true); + expect(acc1?.health).toBe(0.85); + expect(acc1?.cooldownStatus).toContain("1 active (TEST_COOLDOWN)"); + + const acc2 = list.find(a => a.id === "acc2"); + expect(acc2?.enabled).toBe(false); + expect(acc2?.cooldownStatus).toBe("none"); + }); + + test("createAddAccountCallback handles OAuth success", async () => { + const pool = new AccountPool({ accounts: [] }); + const manager = new AccountManager({ pool, client: mockClient }); + + const callback = manager.createAddAccountCallback(); + await callback({ + type: "success", + refresh: "new-refresh", + access: "new-access", + email: "new@test.com" + }); + + expect(pool.count()).toBe(1); + const newAcc = pool.getAccounts()[0]; + expect(newAcc.account.id).toBe("new@test.com"); + expect(newAcc.account.refreshToken).toBe("new-refresh"); + expect(newAcc.auth.access).toBe("new-access"); + + const cached = resolveCachedAuthForAccount("new@test.com"); + expect(cached?.access).toBe("new-access"); + }); + + test("createAddAccountCallback throws on OAuth failure", async () => { + const pool = new AccountPool({ accounts: [] }); + const manager = new AccountManager({ pool, client: mockClient }); + + const callback = manager.createAddAccountCallback(); + + expect(callback({ type: "failed", refresh: "" })).rejects.toThrow("OAuth failed or missing refresh token"); + expect(pool.count()).toBe(0); + }); +}); diff --git a/src/plugin/account-manager.ts b/src/plugin/account-manager.ts new file mode 100644 index 0000000..5e335a5 --- /dev/null +++ b/src/plugin/account-manager.ts @@ -0,0 +1,122 @@ +/** + * Account management utilities for adding, removing, enabling/disabling accounts. + * Integrates with the AccountPool and provides auth method callbacks. + */ + +import type { AccountPool, GeminiAccount, PluginClient } from "./types"; +import { clearCachedAuthForAccount, storeCachedAuthForAccount } from "./cache"; +import { invalidateProjectContextCacheForAccount } from "./project/context"; + +export interface AccountManagerOptions { + pool: AccountPool; + client: PluginClient; + onConfigUpdate?: (accounts: GeminiAccount[]) => void; +} + +export class AccountManager { + private pool: AccountPool; + private client: PluginClient; + private onConfigUpdate?: (accounts: GeminiAccount[]) => void; + + constructor(options: AccountManagerOptions) { + this.pool = options.pool; + this.client = options.client; + this.onConfigUpdate = options.onConfigUpdate; + } + + /** + * Adds a new account to the pool after successful OAuth. + */ + async addAccount(account: GeminiAccount): Promise { + this.pool.addAccount(account); + this.notifyConfigUpdate(); + } + + /** + * Removes an account from the pool, clearing its cache and context. + */ + removeAccount(accountId: string): void { + this.pool.removeAccount(accountId); + clearCachedAuthForAccount(accountId); + invalidateProjectContextCacheForAccount(accountId); + this.notifyConfigUpdate(); + } + + /** + * Toggles an account's enabled state. + */ + toggleAccount(accountId: string, enabled: boolean): void { + this.pool.toggleAccount(accountId, enabled); + this.notifyConfigUpdate(); + } + + /** + * Lists all accounts with their current status. + */ + listAccounts(): { + id: string; + email?: string; + enabled: boolean; + health: number; + cooldownStatus: string; + }[] { + const now = Date.now(); + return this.pool.getAccounts().map((state) => { + const activeCooldowns = Array.from(state.cooldowns.values()).filter((c) => c.expiresAt > now); + return { + id: state.account.id, + email: state.account.email, + enabled: state.account.enabled, + health: state.health.value, + cooldownStatus: activeCooldowns.length > 0 + ? `${activeCooldowns.length} active (${activeCooldowns[0].reason})` + : "none", + }; + }); + } + + /** + * Creates an OAuth callback handler for adding a new account. + * When OAuth succeeds, this adds the account to the pool. + */ + createAddAccountCallback(): (result: { + type: string; + refresh: string; + access?: string; + expires?: number; + email?: string; + }) => Promise { + return async (result) => { + if (result.type !== "success" || !result.refresh) { + throw new Error("OAuth failed or missing refresh token"); + } + const account: GeminiAccount = { + id: result.email || result.refresh.split("|")[0] || `account-${Date.now()}`, + email: result.email, + refreshToken: result.refresh, + enabled: true, + }; + await this.addAccount(account); + // Store initial auth in pool and cache + this.pool.updateAuth(account.id, { + type: "oauth", + refresh: result.refresh, + access: result.access, + expires: result.expires, + }); + storeCachedAuthForAccount(account.id, { + type: "oauth", + refresh: result.refresh, + access: result.access, + expires: result.expires, + }); + }; + } + + private notifyConfigUpdate(): void { + if (this.onConfigUpdate) { + const accounts = this.pool.getAccounts().map((s) => s.account); + this.onConfigUpdate(accounts); + } + } +} diff --git a/src/plugin/account-pool.test.ts b/src/plugin/account-pool.test.ts new file mode 100644 index 0000000..ea0df63 --- /dev/null +++ b/src/plugin/account-pool.test.ts @@ -0,0 +1,240 @@ +import { describe, expect, it } from "bun:test"; + +import { AccountPool } from "./account-pool"; +import type { GeminiAccount, OAuthAuthDetails, ProjectContextResult } from "./types"; + +function createPool(count: number = 2): AccountPool { + const accounts: GeminiAccount[] = Array.from({ length: count }, (_, i) => ({ + id: `acc-${i}`, + email: `acc-${i}@example.com`, + refreshToken: `rt-acc-${i}`, + enabled: true, + })); + return new AccountPool({ accounts, strategy: "health-weighted" }); +} + +describe("AccountPool", () => { + describe("pool construction", () => { + it("initializes with correct account count", () => { + const pool = createPool(2); + expect(pool.count()).toBe(2); + }); + + it("initializes with zero accounts when given an empty list", () => { + const pool = new AccountPool({ accounts: [] }); + expect(pool.count()).toBe(0); + }); + + it("ignores duplicate addAccount calls", () => { + const pool = createPool(1); + pool.addAccount({ id: "acc-0", email: "dup@x.com", refreshToken: "rt-dup", enabled: true }); + expect(pool.count()).toBe(1); + }); + }); + + describe("select", () => { + it("returns undefined on an empty pool", () => { + const pool = new AccountPool({ accounts: [] }); + expect(pool.select()).toBeUndefined(); + }); + + it("returns the only account on a pool of one", () => { + const pool = createPool(1); + const selected = pool.select(); + expect(selected?.account.id).toBe("acc-0"); + }); + + it("returns account with highest health (deterministic if other is zero)", () => { + const pool = createPool(2); + // Set acc-0's health to 0 + pool.updateQuotaRemaining("acc-0", 0); + // We need to make sure health actually goes to 0 or very low + // Actually, HealthTracker compute() might not go to 0 immediately. + // Let's just report many failures. + for (let i = 0; i < 20; i++) pool.reportFailure("acc-0"); + + const selected = pool.select(); + // acc-1 still has perfect health (1.0) + expect(selected?.account.id).toBe("acc-1"); + }); + + it("excludes cooldowned accounts", () => { + const pool = createPool(2); + pool.cooldownAccount("acc-0", "gemini-2.0-flash", 60000, "RATE_LIMIT"); + const selected = pool.select("gemini-2.0-flash"); + expect(selected?.account.id).toBe("acc-1"); + }); + + it("excludes disabled accounts", () => { + const pool = createPool(2); + pool.toggleAccount("acc-0", false); + const selected = pool.select(); + expect(selected?.account.id).toBe("acc-1"); + }); + + it("picks shortest remaining cooldown when all accounts are cooldowned", () => { + const pool = createPool(2); + // Long cooldown on acc-0, short cooldown on acc-1 + pool.cooldownAccount("acc-0", "gemini-2.0-flash", 60000, "RATE_LIMIT"); + pool.cooldownAccount("acc-1", "gemini-2.0-flash", 10000, "RATE_LIMIT"); + const selected = pool.select("gemini-2.0-flash"); + expect(selected?.account.id).toBe("acc-1"); + }); + + it("increments usageCount and updates lastUsed on selection", () => { + const pool = createPool(2); + const selected = pool.select(); + const id = selected?.account.id; + expect(id).toBeDefined(); + const after = pool.getAccount(id!)?.usageCount ?? -1; + expect(after).toBe(1); + }); + }); + + describe("cooldown and health reporting", () => { + it("cooldownAccount sets a cooldown entry (canonicalized)", () => { + const pool = createPool(2); + pool.cooldownAccount("acc-0", "gemini-2.0-flash", 10000, "RATE_LIMIT"); + const cooldowns = pool.getAccount("acc-0")?.cooldowns; + expect(cooldowns?.has("models/gemini-2.0-flash")).toBe(true); + expect(cooldowns?.get("models/gemini-2.0-flash")?.reason).toBe("RATE_LIMIT"); + }); + + it("reportFailure lowers health score", () => { + const pool = createPool(1); + expect(pool.getAccount("acc-0")?.health.value).toBe(1); + pool.reportFailure("acc-0"); + expect(pool.getAccount("acc-0")?.health.value).toBeLessThan(1); + }); + + it("reportResult handles 401 by disabling the account", () => { + const pool = createPool(1); + expect(pool.getAccount("acc-0")?.account.enabled).toBe(true); + pool.reportResult("acc-0", 401); + expect(pool.getAccount("acc-0")?.account.enabled).toBe(false); + }); + + it("reportResult handles success", () => { + const pool = createPool(1); + pool.reportFailure("acc-0"); + const low = pool.getAccount("acc-0")?.health.value ?? 1; + pool.reportResult("acc-0", 200); + const high = pool.getAccount("acc-0")?.health.value ?? 0; + expect(high).toBeGreaterThan(low); + }); + + it("canonicalizes model names in cooldownAccount", () => { + const pool = createPool(1); + pool.cooldownAccount("acc-0", "gemini-1.5-pro", 10000, "RATE_LIMIT"); + const cooldowns = pool.getAccount("acc-0")?.cooldowns; + expect(cooldowns?.has("models/gemini-1.5-pro")).toBe(true); + }); + + it("reportSuccess after failure improves health score", () => { + const pool = createPool(1); + pool.reportFailure("acc-0"); + const afterFailure = pool.getAccount("acc-0")?.health.value ?? 1; + pool.reportSuccess("acc-0", 100); + const afterSuccess = pool.getAccount("acc-0")?.health.value ?? 0; + expect(afterSuccess).toBeGreaterThan(afterFailure); + }); + }); + + describe("withRefreshLock", () => { + it("deduplicates concurrent refreshes for the same account", async () => { + const pool = createPool(2); + let callCount = 0; + const refreshFn = async (): Promise => { + callCount++; + return null; + }; + const p1 = pool.withRefreshLock("acc-0", refreshFn); + const p2 = pool.withRefreshLock("acc-0", refreshFn); + expect(callCount).toBe(1); + expect(p1).toBe(p2); + await p1; + }); + + it("allows concurrent refreshes for different accounts", async () => { + const pool = createPool(2); + let callCount = 0; + const refreshFn = async (): Promise => { + callCount++; + return null; + }; + const p1 = pool.withRefreshLock("acc-0", refreshFn); + const p2 = pool.withRefreshLock("acc-1", refreshFn); + // Different accounts → both refreshFn should be called + expect(p1).not.toBe(p2); + await Promise.all([p1, p2]); + expect(callCount).toBe(2); + }); + }); + + describe("account lifecycle", () => { + it("removeAccount decreases count", () => { + const pool = createPool(3); + expect(pool.count()).toBe(3); + pool.removeAccount("acc-1"); + expect(pool.count()).toBe(2); + expect(pool.getAccount("acc-1")).toBeUndefined(); + }); + + it("toggleAccount enables and disables accounts", () => { + const pool = createPool(2); + expect(pool.getAccount("acc-0")?.account.enabled).toBe(true); + pool.toggleAccount("acc-0", false); + expect(pool.getAccount("acc-0")?.account.enabled).toBe(false); + pool.toggleAccount("acc-0", true); + expect(pool.getAccount("acc-0")?.account.enabled).toBe(true); + }); + + it("fromSingleAccount creates pool-of-one that always returns that account", () => { + const account: GeminiAccount = { + id: "legacy", + email: "legacy@example.com", + refreshToken: "rt-legacy", + enabled: true, + }; + const pool = AccountPool.fromSingleAccount(account); + expect(pool.count()).toBe(1); + const selected = pool.select(); + expect(selected?.account.id).toBe("legacy"); + }); + }); + + describe("auth and project context", () => { + it("updateAuth updates account auth details", () => { + const pool = createPool(2); + const newAuth: OAuthAuthDetails = { + type: "oauth", + refresh: "rt-new", + access: "tok-new", + expires: Date.now() + 3600000, + }; + pool.updateAuth("acc-0", newAuth); + const state = pool.getAccount("acc-0"); + expect(state?.auth.refresh).toBe("rt-new"); + expect(state?.auth.access).toBe("tok-new"); + }); + + it("updateProjectContext updates account project context", () => { + const pool = createPool(2); + const ctx: ProjectContextResult = { + auth: { type: "oauth", refresh: "rt-ctx" }, + effectiveProjectId: "proj-123", + }; + pool.updateProjectContext("acc-0", ctx); + const state = pool.getAccount("acc-0"); + expect(state?.projectContext?.effectiveProjectId).toBe("proj-123"); + expect(state?.projectContext?.auth.refresh).toBe("rt-ctx"); + }); + + it("updateQuotaRemaining updates health quota remaining", () => { + const pool = createPool(1); + expect(pool.getAccount("acc-0")?.health.quotaRemaining).toBe(1); + pool.updateQuotaRemaining("acc-0", 0.3); + expect(pool.getAccount("acc-0")?.health.quotaRemaining).toBe(0.3); + }); + }); +}); diff --git a/src/plugin/account-pool.ts b/src/plugin/account-pool.ts new file mode 100644 index 0000000..820e3b7 --- /dev/null +++ b/src/plugin/account-pool.ts @@ -0,0 +1,188 @@ +/** + * AccountPool manages multiple Gemini accounts with rotation, cooldown tracking, + * and health scoring. It is the source of truth for per-account state. + */ + +import { HealthTracker } from "./health"; +import { selectBestAccount } from "./rotation"; +import { getCanonicalModelName } from "./models"; +import type { + AccountState, + GeminiAccount, + AccountPoolConfig, + OAuthAuthDetails, + ProjectContextResult, + CooldownState, +} from "./types"; + +export class AccountPool { + private states: Map = new Map(); + private trackers: Map = new Map(); + private refreshLocks: Map> = new Map(); + private strategy: string = "health-weighted"; + + constructor(config: AccountPoolConfig) { + this.strategy = config.strategy ?? "health-weighted"; + for (const account of config.accounts) { + this.addAccount(account); + } + } + + addAccount(account: GeminiAccount): void { + if (this.states.has(account.id)) return; + const defaultAuth: OAuthAuthDetails = { + type: "oauth", + refresh: account.refreshToken, + }; + this.states.set(account.id, { + account, + auth: defaultAuth, + cooldowns: new Map(), + lastUsed: 0, + usageCount: 0, + health: { value: 1.0, successRate: 1.0, quotaRemaining: 1.0, latencyScore: 1.0, cooldownScore: 1.0 }, + }); + this.trackers.set(account.id, new HealthTracker()); + } + + removeAccount(accountId: string): void { + this.states.delete(accountId); + this.trackers.delete(accountId); + } + + toggleAccount(accountId: string, enabled: boolean): void { + const state = this.states.get(accountId); + if (state) { + state.account.enabled = enabled; + } + } + + getAccount(accountId: string): AccountState | undefined { + return this.states.get(accountId); + } + + getAccounts(): AccountState[] { + return Array.from(this.states.values()); + } + + count(): number { + return this.states.size; + } + + select(model?: string, url?: string): AccountState | undefined { + const accounts = this.getAccounts(); + const selected = selectBestAccount(accounts, { model, url }); + if (selected) { + selected.lastUsed = Date.now(); + selected.usageCount++; + } + return selected; + } + + cooldownAccount(accountId: string, model: string, durationMs: number, reason: string): void { + const state = this.states.get(accountId); + if (!state) return; + const canonicalModel = getCanonicalModelName(model); + state.cooldowns.set(canonicalModel, { + accountId, + expiresAt: Date.now() + durationMs, + reason, + model: canonicalModel, + }); + const tracker = this.trackers.get(accountId); + if (tracker) { + tracker.recordCooldown(); + state.health = tracker.compute(); + } + } + + reportResult(accountId: string, status: number, latencyMs?: number): void { + if (status >= 200 && status < 300) { + this.reportSuccess(accountId, latencyMs); + } else { + if (status === 401) { + this.toggleAccount(accountId, false); + } + this.reportFailure(accountId); + } + } + + reportSuccess(accountId: string, latencyMs?: number): void { + const state = this.states.get(accountId); + if (!state) return; + const tracker = this.trackers.get(accountId); + if (tracker) { + tracker.recordSuccess(latencyMs); + state.health = tracker.compute(); + } + } + + reportFailure(accountId: string): void { + const state = this.states.get(accountId); + if (!state) return; + const tracker = this.trackers.get(accountId); + if (tracker) { + tracker.recordFailure(); + state.health = tracker.compute(); + } + } + + updateAuth(accountId: string, auth: OAuthAuthDetails): void { + const state = this.states.get(accountId); + if (!state) return; + state.auth = auth; + // Sync the refresh token back to the base account object so it can be persisted + if (auth.refresh && auth.refresh !== state.account.refreshToken) { + state.account.refreshToken = auth.refresh; + } + } + + updateProjectContext(accountId: string, context: ProjectContextResult): void { + const state = this.states.get(accountId); + if (!state) return; + state.projectContext = context; + } + + updateQuotaRemaining(accountId: string, pct: number): void { + const state = this.states.get(accountId); + if (!state) return; + const tracker = this.trackers.get(accountId); + if (tracker) { + tracker.setQuotaRemaining(pct); + state.health = tracker.compute(); + } + } + + getHealthTracker(_accountId: string): HealthTracker | null { + // HealthTracker is internal to AccountPool, not exposed directly + return null; + } + + withRefreshLock( + accountId: string, + refreshFn: () => Promise, + ): Promise { + const existing = this.refreshLocks.get(accountId); + if (existing) return existing; + + const timeout = new Promise((_, reject) => + setTimeout(() => reject(new Error("Refresh lock timeout after 30s")), 30000), + ); + + const lock = Promise.race([refreshFn(), timeout]).finally(() => { + this.refreshLocks.delete(accountId); + }) as Promise; + + this.refreshLocks.set(accountId, lock); + return lock; + } + + /** + * Creates a pool-of-one wrapper for backwards compatibility. + * When no accounts[] are configured, the single legacy credential + * is wrapped in a pool that always returns that account. + */ + static fromSingleAccount(account: GeminiAccount): AccountPool { + return new AccountPool({ accounts: [account], strategy: "health-weighted" }); + } +} diff --git a/src/plugin/cache.ts b/src/plugin/cache.ts index 5b5f8c9..40c11c2 100644 --- a/src/plugin/cache.ts +++ b/src/plugin/cache.ts @@ -63,3 +63,19 @@ export function clearCachedAuth(refresh?: string): void { authCache.delete(key); } } + +// ── Account ID-based cache ──────────────────────────────────────── + +const authCacheByAccountId = new Map(); + +export function resolveCachedAuthForAccount(accountId: string): OAuthAuthDetails | undefined { + return authCacheByAccountId.get(accountId); +} + +export function storeCachedAuthForAccount(accountId: string, auth: OAuthAuthDetails): void { + authCacheByAccountId.set(accountId, auth); +} + +export function clearCachedAuthForAccount(accountId: string): void { + authCacheByAccountId.delete(accountId); +} diff --git a/src/plugin/config-store.ts b/src/plugin/config-store.ts new file mode 100644 index 0000000..87cc526 --- /dev/null +++ b/src/plugin/config-store.ts @@ -0,0 +1,47 @@ +import fs from "node:fs/promises"; +import path from "node:path"; +import os from "node:os"; +import type { GeminiAccount } from "./types"; + +function getConfigFilePath(): string { + const configDir = process.env.OPENCODE_CONFIG_DIR || path.join(os.homedir(), ".config", "opencode"); + const isTest = process.env.NODE_ENV === "test" || process.env.BUN_ENV === "test"; + const fileName = isTest ? "gemini-auth.test.json" : "gemini-auth.json"; + return path.join(configDir, fileName); +} + +export async function loadAccountsFromDisk(): Promise { + try { + const filePath = getConfigFilePath(); + const content = await fs.readFile(filePath, "utf-8"); + const parsed = JSON.parse(content); + if (parsed && Array.isArray(parsed.accounts)) { + return parsed.accounts as GeminiAccount[]; + } + } catch (error) { + // Ignore if file doesn't exist or is invalid + } + return []; +} + +export async function saveAccountsToDisk(accounts: GeminiAccount[]): Promise { + try { + const filePath = getConfigFilePath(); + const dir = path.dirname(filePath); + await fs.mkdir(dir, { recursive: true }); + + const data = { + accounts: accounts.map(a => ({ + id: a.id, + email: a.email, + refreshToken: a.refreshToken, + projectId: a.projectId, + enabled: a.enabled + })) + }; + + await fs.writeFile(filePath, JSON.stringify(data, null, 2), "utf-8"); + } catch (error) { + console.error("[Gemini Auth] Failed to save accounts to disk:", error); + } +} diff --git a/src/plugin/health.test.ts b/src/plugin/health.test.ts new file mode 100644 index 0000000..4c5f740 --- /dev/null +++ b/src/plugin/health.test.ts @@ -0,0 +1,121 @@ +import { describe, expect, it } from "bun:test"; + +import { HealthTracker } from "./health"; + +describe("HealthTracker", () => { + it("starts with a perfect score of 1.0 before any records", () => { + const tracker = new HealthTracker(); + const score = tracker.compute(); + expect(score.value).toBe(1); + expect(score.successRate).toBe(1); + expect(score.quotaRemaining).toBe(1); + // No records yet → no latency penalty → perfect latency score + expect(score.latencyScore).toBe(1); + expect(score.cooldownScore).toBe(1); + }); + + it("maintains a high score after a successful request", () => { + const tracker = new HealthTracker(); + tracker.recordSuccess(100); + const score = tracker.compute(); + // 1 success, 0 failures → successRate = 1.0 + // latency = 100ms → latencyScore = 1 - 100/5000 = 0.98 + // no cooldowns → cooldownScore = 1.0 + // quota = 1.0 (default) + // value = 0.40*1.0 + 0.30*1.0 + 0.15*0.98 + 0.15*1.0 = 0.40+0.30+0.147+0.15 = 0.997 → 1.0 + expect(score.value).toBeGreaterThanOrEqual(0.99); + }); + + it("decreases the score when failures are recorded", () => { + const tracker = new HealthTracker(); + // Start with perfect + const initial = tracker.compute(); + expect(initial.value).toBe(1); + + // Record 3 failures in a row + tracker.recordFailure(); + tracker.recordFailure(); + tracker.recordFailure(); + const afterFailures = tracker.compute(); + // successRate = 0/3 = 0 → weight 40% drops to 0 + // value ≈ 0 + 0.30*1 + 0.15*0.9 + 0.15*1 = 0.585 + expect(afterFailures.value).toBeLessThan(initial.value); + expect(afterFailures.successRate).toBe(0); + }); + + it("decreases cooldown score when cooldowns are recorded", () => { + const tracker = new HealthTracker(); + const before = tracker.compute(); + expect(before.cooldownScore).toBe(1); + + tracker.recordCooldown(); + // Cooldown score decays based on recency; value should drop + const after = tracker.compute(); + expect(after.cooldownScore).toBeLessThan(before.cooldownScore); + }); + + it("reflects quota remaining in the score", () => { + const tracker = new HealthTracker(); + expect(tracker.compute().quotaRemaining).toBe(1); + + tracker.setQuotaRemaining(0.5); + expect(tracker.compute().quotaRemaining).toBe(0.5); + + tracker.setQuotaRemaining(0); + expect(tracker.compute().quotaRemaining).toBe(0); + }); + + it("clamps quota values to [0, 1]", () => { + const tracker = new HealthTracker(); + tracker.setQuotaRemaining(-0.5); + expect(tracker.compute().quotaRemaining).toBe(0); + + tracker.setQuotaRemaining(1.5); + expect(tracker.compute().quotaRemaining).toBe(1); + }); + + it("returns rounded values from compute()", () => { + const tracker = new HealthTracker(); + // Record enough to create non-trivial fractions + tracker.recordSuccess(333); + tracker.recordFailure(); + tracker.recordSuccess(250); + const score = tracker.compute(); + // All values should be rounded to at most 2 decimal places + expect(score.value * 100 % 1).toBe(0); + expect(score.successRate * 100 % 1).toBe(0); + expect(Math.round(score.successRate * 100) / 100).toBe(score.successRate); + }); + + it("resets all state back to defaults", () => { + const tracker = new HealthTracker(); + tracker.recordSuccess(200); + tracker.recordFailure(); + tracker.recordCooldown(); + tracker.setQuotaRemaining(0.3); + + const before = tracker.compute(); + expect(before.value).toBeLessThan(1); + + tracker.reset(); + const after = tracker.compute(); + expect(after.value).toBe(1); + expect(after.successRate).toBe(1); + expect(after.quotaRemaining).toBe(1); + expect(after.cooldownScore).toBe(1); + }); + + it("is O(1): compute() does not mutate internal state", () => { + const tracker = new HealthTracker(); + tracker.recordSuccess(100); + tracker.recordFailure(); + + const first = tracker.compute(); + const second = tracker.compute(); + const third = tracker.compute(); + + // Calling compute() repeatedly without new records returns the same values + expect(first).toEqual(second); + expect(second).toEqual(third); + }); +}); diff --git a/src/plugin/health.ts b/src/plugin/health.ts new file mode 100644 index 0000000..bc0c01c --- /dev/null +++ b/src/plugin/health.ts @@ -0,0 +1,84 @@ +/** + * Tracks per-account health scores with O(1) incremental updates. + * + * Weights: success rate 40%, quota remaining 30%, latency 15%, cooldown frequency 15%. + */ + +import type { HealthScore } from "./types"; + +const WEIGHT_SUCCESS = 0.40; +const WEIGHT_QUOTA = 0.30; +const WEIGHT_LATENCY = 0.15; +const WEIGHT_COOLDOWN = 0.15; + +const DEFAULT_LATENCY_MS = 500; +const MAX_LATENCY_MS = 5000; +const COOLDOWN_PENALTY_WINDOW_MS = 5 * 60 * 1000; // 5 minutes + +export class HealthTracker { + private successCount: number = 0; + private failureCount: number = 0; + private totalLatencyMs: number = 0; + private requestCount: number = 0; + private cooldownCount: number = 0; + private lastCooldownAt: number = 0; + private quotaRemaining: number = 1.0; + + recordSuccess(latencyMs: number = DEFAULT_LATENCY_MS): void { + this.successCount++; + this.requestCount++; + this.totalLatencyMs += latencyMs; + } + + recordFailure(): void { + this.failureCount++; + this.requestCount++; + } + + recordCooldown(): void { + this.cooldownCount++; + this.lastCooldownAt = Date.now(); + } + + setQuotaRemaining(pct: number): void { + this.quotaRemaining = Math.max(0, Math.min(1, pct)); + } + + compute(): HealthScore { + const total = this.successCount + this.failureCount; + const successRate = total > 0 ? this.successCount / total : 1.0; + const avgLatency = this.requestCount > 0 ? this.totalLatencyMs / this.requestCount : 0; + const latencyScore = Math.max(0, 1 - (avgLatency / MAX_LATENCY_MS)); + + // Cooldown score: penalize recent cooldowns + const timeSinceLastCooldown = Date.now() - this.lastCooldownAt; + const cooldownDecay = this.lastCooldownAt > 0 + ? Math.max(0, 1 - (this.cooldownCount * (COOLDOWN_PENALTY_WINDOW_MS / Math.max(timeSinceLastCooldown, 1)))) + : 1.0; + const cooldownScore = Math.max(0, Math.min(1, cooldownDecay)); + + const value = + WEIGHT_SUCCESS * successRate + + WEIGHT_QUOTA * this.quotaRemaining + + WEIGHT_LATENCY * latencyScore + + WEIGHT_COOLDOWN * cooldownScore; + + return { + value: Math.round(value * 100) / 100, + successRate: Math.round(successRate * 100) / 100, + quotaRemaining: Math.round(this.quotaRemaining * 100) / 100, + latencyScore: Math.round(latencyScore * 100) / 100, + cooldownScore: Math.round(cooldownScore * 100) / 100, + }; + } + + reset(): void { + this.successCount = 0; + this.failureCount = 0; + this.totalLatencyMs = 0; + this.requestCount = 0; + this.cooldownCount = 0; + this.lastCooldownAt = 0; + this.quotaRemaining = 1.0; + } +} diff --git a/src/plugin/models.ts b/src/plugin/models.ts new file mode 100644 index 0000000..ada1a03 --- /dev/null +++ b/src/plugin/models.ts @@ -0,0 +1,39 @@ +/** + * Canonical model mapping for Gemini. + * Maps various aliases and formats to a single canonical name to ensure + * cooldowns and quota tracking are consistent across requests. + */ + +const MODEL_MAP: Record = { + "gemini-1.5-pro": "models/gemini-1.5-pro", + "gemini-1.5-flash": "models/gemini-1.5-flash", + "gemini-2.0-flash": "models/gemini-2.0-flash", + "gemini-2.0-flash-lite": "models/gemini-2.0-flash-lite", + "gemini-2.0-pro-exp-02-05": "models/gemini-2.0-pro-exp-02-05", +}; + +/** + * Returns the canonical name for a model. + * If no mapping exists, returns the original name. + */ +export function getCanonicalModelName(model: string): string { + // If it already starts with models/, use it as is + if (model.startsWith("models/")) { + return model; + } + + // Check the map + const mapped = MODEL_MAP[model]; + if (mapped) { + return mapped; + } + + // Fallback: if it's a versioned model without prefix (e.g. gemini-1.5-pro-002) + if (model.startsWith("gemini-")) { + // If it has a specific version suffix, we might still want to canonicalize the base + // but for now let's just prefix it if it's a known gemini model. + return `models/${model}`; + } + + return model; +} diff --git a/src/plugin/oauth-authorize.ts b/src/plugin/oauth-authorize.ts index c00733c..ea53ef0 100644 --- a/src/plugin/oauth-authorize.ts +++ b/src/plugin/oauth-authorize.ts @@ -6,7 +6,8 @@ import { isGeminiDebugEnabled, logGeminiDebugMessage } from "./debug"; import { resolveProjectContextFromAccessToken } from "./project"; import { resolveConfiguredProjectId } from "./provider"; import { startOAuthListener, type OAuthListener } from "./server"; -import type { OAuthAuthDetails } from "./types"; +import type { OAuthAuthDetails, AuthMethod } from "./types"; +import type { AccountManager } from "./account-manager"; /** * Builds the OAuth authorize callback used by plugin auth methods. @@ -219,6 +220,80 @@ function openBrowserUrl(url: string): void { } catch {} } +/** + * Creates an OAuth authorize method scoped to adding a new account. + * Similar to createOAuthAuthorizeMethod but returns the account info + * for pool integration instead of replacing the global auth. + */ +export function createOAuthAuthorizeMethodForAccount(options?: { + onAccountAdded?: (account: { email?: string; refreshToken: string; access?: string; expires?: number }) => void; +}): () => Promise<{ + url: string; + instructions: string; + method: string; + callback: (() => Promise) | ((callbackUrl: string) => Promise); +}> { + return async () => { + const baseMethod = createOAuthAuthorizeMethod(); + const result = await baseMethod(); + + return { + ...result, + callback: async (callbackUrl?: string) => { + const exchangeResult = typeof result.callback === "function" && result.callback.length > 0 + ? await (result.callback as any)(callbackUrl) + : await (result.callback as any)(); + + if (exchangeResult.type === "success" && options?.onAccountAdded) { + options.onAccountAdded({ + email: exchangeResult.email, + refreshToken: exchangeResult.refresh, + access: exchangeResult.access, + expires: exchangeResult.expires + }); + } + return exchangeResult; + } + }; + }; +} + +/** + * Creates an "Add Gemini Account" auth method that integrates with AccountManager. + */ +export function createAddAccountAuthMethod(getAccountManager: () => AccountManager): AuthMethod { + return { + label: "Add Gemini Account", + type: "oauth", + authorize: async () => { + const baseMethod = createOAuthAuthorizeMethod(); + const baseResult = await baseMethod(); + + return { + ...baseResult, + callback: async (callbackUrl?: string) => { + const exchangeResult = typeof baseResult.callback === "function" && baseResult.callback.length > 0 + ? await (baseResult.callback as any)(callbackUrl) + : await (baseResult.callback as any)(); + + if (exchangeResult.type === "success") { + const accountManager = getAccountManager(); + const addCallback = accountManager.createAddAccountCallback(); + await addCallback({ + type: "success", + refresh: exchangeResult.refresh, + access: exchangeResult.access, + expires: exchangeResult.expires, + email: exchangeResult.email, + }); + } + return exchangeResult as GeminiTokenExchangeResult; + }, + }; + }, + }; +} + function shouldIgnoreMalformedAuthCode(result: GeminiTokenExchangeResult): boolean { if (result.type !== "failed") { return false; diff --git a/src/plugin/plugin.test.ts b/src/plugin/plugin.test.ts new file mode 100644 index 0000000..5cec208 --- /dev/null +++ b/src/plugin/plugin.test.ts @@ -0,0 +1,128 @@ +import { afterEach, describe, expect, it, mock } from "bun:test"; + +mock.module("./config-store", () => ({ + loadAccountsFromDisk: async () => [], + saveAccountsToDisk: async () => {} +})); + +import { GeminiCLIOAuthPlugin, getLatestGeminiPool, resetPluginState } from "../plugin"; +import type { PluginClient, Provider } from "./types"; + +function makeMockClient(): PluginClient { + return { + auth: { set: async () => {} }, + config: { + get: async () => ({ data: undefined }), + }, + tui: { + showToast: async () => ({}), + }, + }; +} + +function makeOAuthAuth() { + return { + type: "oauth" as const, + refresh: "test-refresh-token|test-project", + access: "test-access-token", + expires: Date.now() + 3600000, + }; +} + +describe("multi-account plugin loader", () => { + afterEach(() => { + resetPluginState(); + }); + + it("creates pool-of-one when no accounts configured", async () => { + const plugin = await GeminiCLIOAuthPlugin({ client: makeMockClient() }); + const getAuth = mock(async () => makeOAuthAuth()); + const provider: Provider = { models: {} }; + + const result = await plugin.auth.loader(getAuth, provider); + expect(result).not.toBeNull(); + + const pool = getLatestGeminiPool(); + expect(pool).toBeDefined(); + expect(pool!.count()).toBe(1); + }); + + it("creates pool with accounts from provider options", async () => { + const plugin = await GeminiCLIOAuthPlugin({ client: makeMockClient() }); + const getAuth = mock(async () => makeOAuthAuth()); + const provider: Provider = { + models: {}, + options: { + accounts: [ + { id: "user1", refreshToken: "rt1", enabled: true }, + { id: "user2", refreshToken: "rt2", enabled: true }, + { id: "user3", refreshToken: "rt3", enabled: true }, + ], + }, + }; + + const result = await plugin.auth.loader(getAuth, provider); + expect(result).not.toBeNull(); + + const pool = getLatestGeminiPool(); + expect(pool).toBeDefined(); + expect(pool!.count()).toBe(3); + }); + + it("returns loader result with fetch function", async () => { + const plugin = await GeminiCLIOAuthPlugin({ client: makeMockClient() }); + const getAuth = mock(async () => makeOAuthAuth()); + const provider: Provider = { models: {} }; + + const result = await plugin.auth.loader(getAuth, provider); + expect(result).not.toBeNull(); + expect(typeof result!.fetch).toBe("function"); + }); + + it("returns null for non-OAuth auth", async () => { + const plugin = await GeminiCLIOAuthPlugin({ client: makeMockClient() }); + const getAuth = mock(async () => ({ type: "api" })); + const provider: Provider = { models: {} }; + + const result = await plugin.auth.loader(getAuth, provider); + expect(result).toBeNull(); + }); + + it("updates auth for matching account from single getAuth() result", async () => { + const plugin = await GeminiCLIOAuthPlugin({ client: makeMockClient() }); + const auth = makeOAuthAuth(); + const getAuth = mock(async () => auth); + const provider: Provider = { + models: {}, + options: { + accounts: [ + { id: "user1", refreshToken: auth.refresh, enabled: true }, + { id: "user2", refreshToken: "rt2", enabled: true }, + ], + }, + }; + + await plugin.auth.loader(getAuth, provider); + const pool = getLatestGeminiPool(); + expect(pool).toBeDefined(); + + const acct1 = pool!.getAccount("user1"); + const acct2 = pool!.getAccount("user2"); + expect(acct1?.auth.access).toBe("test-access-token"); + expect(acct2?.auth.access).toBeUndefined(); + }); + + it("backwards compatible: single-account without accounts[] config works", async () => { + const plugin = await GeminiCLIOAuthPlugin({ client: makeMockClient() }); + const auth = makeOAuthAuth(); + const getAuth = mock(async () => auth); + const provider: Provider = { models: {} }; + + const result = await plugin.auth.loader(getAuth, provider); + expect(result).not.toBeNull(); + + const pool = getLatestGeminiPool(); + expect(pool).toBeDefined(); + expect(pool!.count()).toBe(1); + }); +}); diff --git a/src/plugin/project/context.test.ts b/src/plugin/project/context.test.ts new file mode 100644 index 0000000..698b3d9 --- /dev/null +++ b/src/plugin/project/context.test.ts @@ -0,0 +1,176 @@ +import { beforeEach, describe, expect, it, mock } from "bun:test"; + +import { + ensureProjectContextForAccount, + invalidateProjectContextCacheForAccount, +} from "../project"; +import type { OAuthAuthDetails, PluginClient } from "../types"; +import { buildProjectCacheKeyForAccount } from "./utils"; + +const baseAuth: OAuthAuthDetails = { + type: "oauth", + refresh: "refresh-token", + access: "access-token", + expires: Date.now() + 60_000, +}; + +function createClient(): PluginClient { + return { + auth: { + set: mock(async () => {}), + }, + } as PluginClient; +} + +describe("buildProjectCacheKeyForAccount", () => { + it("returns a stable key for an account and project", () => { + const key = buildProjectCacheKeyForAccount("alice@test.com", "my-project"); + expect(key).toBe("account:alice@test.com|project:my-project"); + }); + + it("returns a stable key with default project when projectId is omitted", () => { + const key = buildProjectCacheKeyForAccount("bob@test.com"); + expect(key).toBe("account:bob@test.com|project:default"); + }); +}); + +describe("ensureProjectContextForAccount", () => { + beforeEach(() => { + mock.restore(); + }); + + it("caches by account ID and returns the same result on repeated calls", async () => { + let fetchCount = 0; + const fetchMock = mock(async (input: RequestInfo) => { + const url = typeof input === "string" ? input : input.toString(); + if (url.includes(":loadCodeAssist")) { + fetchCount++; + return new Response( + JSON.stringify({ + currentTier: { id: "free-tier" }, + cloudaicompanionProject: "projects/server-project", + }), + { status: 200 }, + ); + } + throw new Error(`Unexpected fetch to ${url}`); + }); + (globalThis as { fetch: typeof fetch }).fetch = fetchMock as unknown as typeof fetch; + + const client = createClient(); + + const first = await ensureProjectContextForAccount( + baseAuth, + client, + "alice@test.com", + undefined, + "gemini-3-flash-preview", + ); + + const second = await ensureProjectContextForAccount( + baseAuth, + client, + "alice@test.com", + undefined, + "gemini-3-flash-preview", + ); + + // Second call should hit cache, so fetchCount === 1 + expect(fetchCount).toBe(1); + expect(first.effectiveProjectId).toBe("projects/server-project"); + expect(second.effectiveProjectId).toBe(first.effectiveProjectId); + }); + + it("maintains separate cache entries for different accounts with the same project", async () => { + let fetchCount = 0; + const fetchMock = mock(async (input: RequestInfo) => { + const url = typeof input === "string" ? input : input.toString(); + if (url.includes(":loadCodeAssist")) { + fetchCount++; + return new Response( + JSON.stringify({ + currentTier: { id: "free-tier" }, + cloudaicompanionProject: "projects/server-project", + }), + { status: 200 }, + ); + } + throw new Error(`Unexpected fetch to ${url}`); + }); + (globalThis as { fetch: typeof fetch }).fetch = fetchMock as unknown as typeof fetch; + + const client = createClient(); + + const aliceResult = await ensureProjectContextForAccount( + baseAuth, + client, + "alice@test.com", + "shared-project", + "gemini-3-flash-preview", + ); + + const bobResult = await ensureProjectContextForAccount( + baseAuth, + client, + "bob@test.com", + "shared-project", + "gemini-3-flash-preview", + ); + + // Each account should have triggered a separate fetch (different cache keys) + expect(fetchCount).toBe(2); + expect(aliceResult.effectiveProjectId).toBeTruthy(); + expect(bobResult.effectiveProjectId).toBeTruthy(); + }); +}); + +describe("invalidateProjectContextCacheForAccount", () => { + beforeEach(() => { + mock.restore(); + }); + + it("removes the correct cache entry, forcing a new fetch", async () => { + let fetchCount = 0; + const fetchMock = mock(async (input: RequestInfo) => { + const url = typeof input === "string" ? input : input.toString(); + if (url.includes(":loadCodeAssist")) { + fetchCount++; + return new Response( + JSON.stringify({ + currentTier: { id: "free-tier" }, + cloudaicompanionProject: "projects/server-project", + }), + { status: 200 }, + ); + } + throw new Error(`Unexpected fetch to ${url}`); + }); + (globalThis as { fetch: typeof fetch }).fetch = fetchMock as unknown as typeof fetch; + + const client = createClient(); + + // First call populates the cache + await ensureProjectContextForAccount( + baseAuth, + client, + "alice@test.com", + "my-project", + "gemini-3-flash-preview", + ); + + // Invalidate + invalidateProjectContextCacheForAccount("alice@test.com", "my-project"); + + // Second call should fetch again + await ensureProjectContextForAccount( + baseAuth, + client, + "alice@test.com", + "my-project", + "gemini-3-flash-preview", + ); + + // Two fetches because invalidation removed the cache entry + expect(fetchCount).toBe(2); + }); +}); diff --git a/src/plugin/project/context.ts b/src/plugin/project/context.ts index 6c51bc5..92155d9 100644 --- a/src/plugin/project/context.ts +++ b/src/plugin/project/context.ts @@ -5,6 +5,7 @@ import { loadManagedProject, onboardManagedProject } from "./api"; import { FREE_TIER_ID, LEGACY_TIER_ID, ProjectIdRequiredError } from "./types"; import { buildIneligibleTierMessage, + buildProjectCacheKeyForAccount, getCacheKey, normalizeProjectId, pickOnboardTier, @@ -204,3 +205,30 @@ function buildProjectCacheKey(auth: OAuthAuthDetails, configuredProjectId?: stri const project = configuredProjectId?.trim() ?? ""; return project ? `${base}|cfg:${project}` : base; } + +/** + * Ensures project context for a specific account, caching by account ID. + */ +export async function ensureProjectContextForAccount( + auth: OAuthAuthDetails, + client: PluginClient, + accountId: string, + configuredProjectId?: string, + userAgentModel?: string, +): Promise { + const cacheKey = buildProjectCacheKeyForAccount(accountId, configuredProjectId); + const cached = projectContextResultCache.get(cacheKey); + if (cached) return cached; + + const context = await ensureProjectContext(auth, client, configuredProjectId, userAgentModel); + projectContextResultCache.set(cacheKey, context); + return context; +} + +/** + * Invalidates project context cache for a specific account. + */ +export function invalidateProjectContextCacheForAccount(accountId: string, projectId?: string): void { + const cacheKey = buildProjectCacheKeyForAccount(accountId, projectId); + projectContextResultCache.delete(cacheKey); +} diff --git a/src/plugin/project/index.ts b/src/plugin/project/index.ts index 6edd88b..303575a 100644 --- a/src/plugin/project/index.ts +++ b/src/plugin/project/index.ts @@ -1,6 +1,8 @@ export { loadManagedProject, onboardManagedProject, retrieveUserQuota } from "./api"; export { ensureProjectContext, + ensureProjectContextForAccount, invalidateProjectContextCache, + invalidateProjectContextCacheForAccount, resolveProjectContextFromAccessToken, } from "./context"; diff --git a/src/plugin/project/utils.ts b/src/plugin/project/utils.ts index 341b4d3..9b1223c 100644 --- a/src/plugin/project/utils.ts +++ b/src/plugin/project/utils.ts @@ -139,3 +139,11 @@ export function getCacheKey(auth: OAuthAuthDetails): string | undefined { const refresh = auth.refresh?.trim(); return refresh ? refresh : undefined; } + +/** + * Builds a project context cache key for a specific account. + * Uses account ID (email) instead of packed refresh string. + */ +export function buildProjectCacheKeyForAccount(accountId: string, projectId?: string): string { + return `account:${accountId}|project:${projectId ?? "default"}`; +} diff --git a/src/plugin/provider.ts b/src/plugin/provider.ts index e175a96..05374f1 100644 --- a/src/plugin/provider.ts +++ b/src/plugin/provider.ts @@ -1,7 +1,11 @@ import type { Config } from "@opencode-ai/sdk"; import { GEMINI_PROVIDER_ID } from "../constants"; -import type { PluginClient, Provider } from "./types"; +import type { + GeminiAccount, + PluginClient, + Provider, +} from "./types"; interface ResolveConfiguredProjectIdInput { provider?: Provider | null; @@ -68,3 +72,38 @@ function normalizeProjectId(value: unknown): string | undefined { const trimmed = value.trim(); return trimmed || undefined; } + +/** + * Resolves project ID for a specific account, using the account's optional override. + */ +export function resolveProjectIdForAccount( + account: GeminiAccount, + fallbackProjectId?: string, +): string | undefined { + return account.projectId ?? fallbackProjectId; +} + +/** + * Reads accounts[] from the provider's options. + * Returns undefined when no accounts are configured (pool-of-one fallback). + */ +export function resolveAccountsFromProvider(provider?: Provider | null): GeminiAccount[] | undefined { + if (!provider || typeof provider !== "object") { + return undefined; + } + const options = provider.options; + if (!options || typeof options !== "object") { + return undefined; + } + const raw = (options as Record).accounts; + if (!Array.isArray(raw) || raw.length === 0) { + return undefined; + } + const accounts = raw.filter( + (a): a is GeminiAccount => + a !== null && + typeof a === "object" && + typeof (a as GeminiAccount).id === "string", + ); + return accounts.length > 0 ? accounts : undefined; +} diff --git a/src/plugin/quota.test.ts b/src/plugin/quota.test.ts index 13b17d5..b32f018 100644 --- a/src/plugin/quota.test.ts +++ b/src/plugin/quota.test.ts @@ -1,5 +1,6 @@ import { afterEach, beforeEach, describe, expect, it } from "bun:test"; -import { formatGeminiQuotaOutput, formatRelativeResetTime } from "./quota"; +import { formatAggregateSummary, formatGeminiQuotaOutput, formatRelativeResetTime } from "./quota"; +import type { AccountState } from "./types"; import type { RetrieveUserQuotaBucket } from "./project/types"; const REAL_DATE_NOW = Date.now; @@ -69,6 +70,7 @@ describe("formatGeminiQuotaOutput", () => { const output = formatGeminiQuotaOutput("test-project", buckets); expect(output).toContain("Gemini quota usage for project `test-project`"); expect(output).toContain("Variant"); + expect(output).toContain("Used"); expect(output).toContain("Remaining"); expect(output).toContain("Reset"); expect(output).not.toContain("Type"); @@ -82,6 +84,12 @@ describe("formatGeminiQuotaOutput", () => { expect(output).toContain( "▓▓▓▓▓▓▓▓▓▓░░░░░░░░░░ 50.0% (100 left)", ); + // Verify used column shows calculated values when fraction + amount are available + expect(output).toContain("60"); // 200 total - 140 remaining = 60 used (70% remaining) + expect(output).toContain("100"); // 200 total - 100 remaining = 100 used (50% remaining) + // Verify used column shows percentage when only fraction is available + expect(output).toContain("20.0%"); // 100% - 80% remaining = 20% used + expect(output).toContain("5.0%"); // 100% - 95% remaining = 5% used expect(output.indexOf("Gemini 3 (1 model, 1 bucket)")).toBeLessThan( output.indexOf("Gemini 2.5 (2 models, 3 buckets)"), ); @@ -113,9 +121,82 @@ describe("formatGeminiQuotaOutput", () => { const output = formatGeminiQuotaOutput("test-project", buckets); expect(output).toContain("Type"); + expect(output).toContain("Used"); + expect(output).toContain("10.0%"); // 100% - 90% = 10% used + expect(output).toContain("20.0%"); // 100% - 80% = 20% used expect(output).toContain("Gemini 2.5 (1 model, 2 buckets)"); expect(output).toContain("REQUESTS"); expect(output).toContain("TOKENS"); expect(output).toContain(" ↳ vertex"); }); }); + +describe("formatAggregateSummary", () => { + it("returns summary for single account", () => { + const accounts: AccountState[] = [{ + account: { id: "user1", refreshToken: "rt1", enabled: true }, + auth: { type: "oauth", refresh: "rt1" }, + cooldowns: new Map(), + lastUsed: Date.now(), + usageCount: 0, + health: { value: 1.0, successRate: 1.0, quotaRemaining: 0.5, latencyScore: 1.0, cooldownScore: 1.0 }, + }]; + + const result = formatAggregateSummary(accounts); + expect(result).toContain("Aggregate:"); + expect(result).toContain("50.0%"); + expect(result).toContain("1 account(s)"); + }); + + it("computes average across multiple accounts", () => { + const now = Date.now(); + const accounts: AccountState[] = [ + { + account: { id: "user1", refreshToken: "rt1", enabled: true }, + auth: { type: "oauth", refresh: "rt1" }, + cooldowns: new Map(), + lastUsed: now, + usageCount: 0, + health: { value: 0.8, successRate: 0.8, quotaRemaining: 0.9, latencyScore: 1.0, cooldownScore: 1.0 }, + }, + { + account: { id: "user2", refreshToken: "rt2", enabled: true, email: "user2@test.com" }, + auth: { type: "oauth", refresh: "rt2" }, + cooldowns: new Map(), + lastUsed: now, + usageCount: 0, + health: { value: 0.4, successRate: 0.4, quotaRemaining: 0.1, latencyScore: 0.5, cooldownScore: 0.5 }, + }, + ]; + + const result = formatAggregateSummary(accounts); + // (0.9 + 0.1) / 2 = 0.5 = 50.0% + expect(result).toContain("50.0%"); + expect(result).toContain("2 account(s)"); + }); + + it("returns 0% when all accounts have zero remaining quota", () => { + const now = Date.now(); + const accounts: AccountState[] = [ + { + account: { id: "user1", refreshToken: "rt1", enabled: true }, + auth: { type: "oauth", refresh: "rt1" }, + cooldowns: new Map(), + lastUsed: now, + usageCount: 0, + health: { value: 0.3, successRate: 0.5, quotaRemaining: 0.0, latencyScore: 1.0, cooldownScore: 1.0 }, + }, + { + account: { id: "user2", refreshToken: "rt2", enabled: true }, + auth: { type: "oauth", refresh: "rt2" }, + cooldowns: new Map(), + lastUsed: now, + usageCount: 0, + health: { value: 0.3, successRate: 0.5, quotaRemaining: 0.0, latencyScore: 1.0, cooldownScore: 1.0 }, + }, + ]; + + const result = formatAggregateSummary(accounts); + expect(result).toContain("0.0%"); + }); +}); diff --git a/src/plugin/quota.ts b/src/plugin/quota.ts index 1333eeb..f298d0f 100644 --- a/src/plugin/quota.ts +++ b/src/plugin/quota.ts @@ -1,10 +1,11 @@ import { tool } from "@opencode-ai/plugin"; import { accessTokenExpired, isOAuthAuth } from "./auth"; import { resolveCachedAuth } from "./cache"; -import { ensureProjectContext, retrieveUserQuota } from "./project"; +import { ensureProjectContextForAccount, retrieveUserQuota } from "./project"; import type { RetrieveUserQuotaBucket } from "./project/types"; -import { refreshAccessToken } from "./token"; -import type { GetAuth, PluginClient } from "./types"; +import { refreshAccessToken, refreshAccessTokenForAccount } from "./token"; +import type { AccountPool } from "./account-pool"; +import type { AccountState, GetAuth, PluginClient } from "./types"; export const GEMINI_QUOTA_TOOL_NAME = "gemini_quota"; @@ -13,6 +14,7 @@ interface GeminiQuotaToolDependencies { getAuthResolver: () => GetAuth | undefined; getConfiguredProjectId: () => string | undefined; getUserAgentModel: () => string | undefined; + getPool?: () => AccountPool | undefined; } export function createGeminiQuotaTool({ @@ -20,6 +22,7 @@ export function createGeminiQuotaTool({ getAuthResolver, getConfiguredProjectId, getUserAgentModel, + getPool, }: GeminiQuotaToolDependencies) { return tool({ description: @@ -50,9 +53,19 @@ export function createGeminiQuotaTool({ } try { - const projectContext = await ensureProjectContext( + const pool = getPool?.(); + if (pool && pool.count() > 1) { + return await formatMultiAccountQuotaOutput( + pool, + client, + getUserAgentModel(), + ); + } + + const projectContext = await ensureProjectContextForAccount( authRecord, client, + "default", getConfiguredProjectId(), getUserAgentModel(), ); @@ -102,8 +115,8 @@ export function formatGeminiQuotaOutput( `Gemini quota usage for project \`${projectId}\``, "", showTokenType - ? ` ↳ ${pad("Variant", variantWidth)} Remaining Reset Type` - : ` ↳ ${pad("Variant", variantWidth)} Remaining Reset`, + ? ` ↳ ${pad("Variant", variantWidth)} Used Remaining Reset Type` + : ` ↳ ${pad("Variant", variantWidth)} Used Remaining Reset`, ]; for (let index = 0; index < versionGroups.length; index += 1) { @@ -120,8 +133,8 @@ export function formatGeminiQuotaOutput( for (const row of model.rows) { lines.push( showTokenType - ? ` ↳ ${pad(row.variant, variantWidth)} ${pad(row.usageRemaining, 27)} ${pad(row.resetValue, 8)} ${row.tokenType}` - : ` ↳ ${pad(row.variant, variantWidth)} ${pad(row.usageRemaining, 27)} ${row.resetValue}`, + ? ` ↳ ${pad(row.variant, variantWidth)} ${pad(row.usageUsed, 11)} ${pad(row.usageRemaining, 27)} ${pad(row.resetValue, 8)} ${row.tokenType}` + : ` ↳ ${pad(row.variant, variantWidth)} ${pad(row.usageUsed, 11)} ${pad(row.usageRemaining, 27)} ${row.resetValue}`, ); } } @@ -149,26 +162,43 @@ function compareQuotaBuckets( return (left.resetTime ?? "").localeCompare(right.resetTime ?? ""); } -function formatUsageRemaining(bucket: RetrieveUserQuotaBucket): string { +interface UsageInfo { + used: string; + remaining: string; +} + +function formatUsageInfo(bucket: RetrieveUserQuotaBucket): UsageInfo { const remainingAmount = formatRemainingAmount(bucket.remainingAmount); const remainingFraction = bucket.remainingFraction; const hasFraction = typeof remainingFraction === "number" && Number.isFinite(remainingFraction); + if (hasFraction && remainingAmount) { + const parsed = Number.parseInt(bucket.remainingAmount!, 10); + if (Number.isFinite(parsed) && remainingFraction > 0 && remainingFraction <= 1) { + const total = Math.round(parsed / remainingFraction); + const used = total - parsed; + return { + used: used.toLocaleString("en-US"), + remaining: `${buildProgressBar(remainingFraction)} ${(remainingFraction * 100).toFixed(1)}% (${remainingAmount} left)`, + }; + } + } + if (hasFraction) { const clamped = clamp(remainingFraction, 0, 1); - const percent = (clamped * 100).toFixed(1); - const bar = buildProgressBar(clamped); - return remainingAmount - ? `${bar} ${percent}% (${remainingAmount} left)` - : `${bar} ${percent}%`; + const usedPercent = ((1 - clamped) * 100).toFixed(1); + return { + used: `${usedPercent}%`, + remaining: `${buildProgressBar(clamped)} ${(clamped * 100).toFixed(1)}%`, + }; } if (remainingAmount) { - return remainingAmount; + return { used: "unknown", remaining: remainingAmount }; } - return "unknown"; + return { used: "unknown", remaining: "unknown" }; } function formatRemainingAmount(value: string | undefined): string | undefined { @@ -243,6 +273,7 @@ function normalizeTokenType(bucket: RetrieveUserQuotaBucket): string { interface GroupedQuotaRow { variant: string; + usageUsed: string; usageRemaining: string; resetValue: string; tokenType: string; @@ -260,7 +291,7 @@ function groupQuotaRows(sortedBuckets: RetrieveUserQuotaBucket[]): GroupedQuotaM for (const bucket of sortedBuckets) { const modelId = bucket.modelId?.trim() || "unknown-model"; const { baseModel, variant } = splitModelVariant(modelId); - const usageRemaining = formatUsageRemaining(bucket); + const usageInfo = formatUsageInfo(bucket); const resetLabel = formatRelativeResetTime(bucket.resetTime); const resetValue = resetLabel?.replace("resets in ", "") ?? "-"; const tokenType = normalizeTokenType(bucket); @@ -269,7 +300,8 @@ function groupQuotaRows(sortedBuckets: RetrieveUserQuotaBucket[]): GroupedQuotaM if (existing) { existing.rows.push({ variant, - usageRemaining, + usageUsed: usageInfo.used, + usageRemaining: usageInfo.remaining, resetValue, tokenType, }); @@ -281,7 +313,8 @@ function groupQuotaRows(sortedBuckets: RetrieveUserQuotaBucket[]): GroupedQuotaM version: extractModelVersion(baseModel), rows: [{ variant, - usageRemaining, + usageUsed: usageInfo.used, + usageRemaining: usageInfo.remaining, resetValue, tokenType, }], @@ -384,3 +417,124 @@ function splitModelVariant(modelId: string): { baseModel: string; variant: strin variant: "default", }; } + +/** + * Formats a multi-account quota output with per-account sections and an aggregate summary. + */ +async function formatMultiAccountQuotaOutput( + pool: AccountPool, + client: PluginClient, + userAgentModel?: string, +): Promise { + const accounts = pool.getAccounts(); + const lines: string[] = ["Gemini quota usage (all accounts):\n"]; + + // Track total remaining fraction to compute an accurate aggregate + let sumOfFractions = 0; + let fractionCount = 0; + + for (const state of accounts) { + let accountLabel = state.account.email ?? state.account.id; + if (accountLabel.length > 30 && !state.account.email) { + accountLabel = `${accountLabel.slice(0, 10)}...${accountLabel.slice(-5)}`; + } + + lines.push(`Account: ${accountLabel}`); + let accessToken = state.auth.access; + + // Attempt to refresh or retrieve access token if missing/expired and we have a refresh token + if ((!accessToken || accessTokenExpired(state.auth)) && state.auth.refresh) { + try { + const refreshed = await refreshAccessTokenForAccount(pool, state.account.id, client); + if (refreshed && refreshed.access) { + accessToken = refreshed.access; + } + } catch (error) { + // Fall back to missing token logic + } + } + + if (!accessToken) { + lines.push(" (no access token available)"); + lines.push(""); + continue; + } + + try { + const projectContext = await ensureProjectContextForAccount( + state.auth, + client, + state.account.id, + state.account.projectId, + userAgentModel, + ); + + const projectId = projectContext.effectiveProjectId; + if (!projectId) { + lines.push(" Quota lookup failed: no Google Cloud project could be resolved."); + lines.push(""); + continue; + } + + const quota = await retrieveUserQuota(accessToken, projectId, userAgentModel); + if (!quota?.buckets?.length) { + lines.push(` No quota buckets returned for project \`${projectId}\`.`); + } else { + // Compute an average remaining quota for THIS account to update the pool health + let accRemaining = 0; + let accBucketCount = 0; + + for (const bucket of quota.buckets) { + if (typeof bucket.remainingFraction === "number") { + accRemaining += bucket.remainingFraction; + accBucketCount++; + + sumOfFractions += bucket.remainingFraction; + fractionCount++; + } + } + + if (accBucketCount > 0) { + pool.updateQuotaRemaining(state.account.id, accRemaining / accBucketCount); + } + + const formatted = formatGeminiQuotaOutput(projectId, quota.buckets); + const quotaLines = formatted.split("\n"); + for (let i = 0; i < quotaLines.length; i++) { + if (i === 0) continue; + if (quotaLines[i] === "") { + if (i === 1) continue; + lines.push(""); + } else { + lines.push(` ${quotaLines[i]}`); + } + } + } + } catch (error) { + const message = error instanceof Error ? error.message : "unknown error"; + lines.push(` Quota lookup failed: ${message}`); + } + lines.push(""); + } + + const avgRemaining = fractionCount > 0 ? (sumOfFractions / fractionCount) * 100 : 0; + lines.push(`Aggregate: ${avgRemaining.toFixed(1)}% average remaining across ${accounts.length} account(s)`); + + return lines.join("\n"); +} + +/** + * Computes the average quota remaining across all accounts. + * Uses health.quotaRemaining as a proxy for per-account remaining quota percentage. + */ +export function formatAggregateSummary(accounts: AccountState[]): string { + // This is kept for backwards compatibility in tests, but in production + // formatMultiAccountQuotaOutput handles the aggregate directly now + let totalRemaining = 0; + const totalCapacity = accounts.length; + for (const account of accounts) { + totalRemaining += account.health.quotaRemaining; + } + const avgRemaining = totalCapacity > 0 ? (totalRemaining / totalCapacity) * 100 : 0; + return `Aggregate: ${avgRemaining.toFixed(1)}% average remaining across ${totalCapacity} account(s)`; +} diff --git a/src/plugin/retry.test.ts b/src/plugin/retry.test.ts index 0aa173f..c67974a 100644 --- a/src/plugin/retry.test.ts +++ b/src/plugin/retry.test.ts @@ -1,7 +1,8 @@ import { afterEach, beforeEach, describe, expect, it, mock } from "bun:test"; -import { fetchWithRetry, retryInternals } from "./retry"; +import { buildRetryThrottleKey, fetchWithRetry, formatAllAccountsExhaustedMessage, retryInternals } from "./retry"; import { classifyQuotaResponse } from "./retry/quota"; +import type { AccountState, CooldownState } from "./types"; const originalSetTimeout = globalThis.setTimeout; const scheduledDelays: number[] = []; @@ -295,6 +296,178 @@ describe("fetchWithRetry", () => { expect(scheduledDelays[1]).toBeGreaterThan(0); expect(scheduledDelays[1]).toBeLessThanOrEqual(1500); }); + + it("accepts accountId parameter without breaking existing behavior", async () => { + const fetchMock = mock(async () => { + if (fetchMock.mock.calls.length === 1) { + const err = new Error("socket reset") as Error & { code?: string }; + err.code = "ECONNRESET"; + throw err; + } + return new Response("ok", { status: 200 }); + }); + (globalThis as { fetch: typeof fetch }).fetch = fetchMock as unknown as typeof fetch; + + const response = await fetchWithRetry( + "https://example.com", + { method: "POST", body: JSON.stringify({ hello: "world" }) }, + "test-account", + ); + + expect(response.status).toBe(200); + expect(fetchMock.mock.calls.length).toBe(2); + }); + + it("includes terminal 429 headers when accountId is provided", async () => { + const fetchMock = mock(async () => makeQuota429("QUOTA_EXHAUSTED")); + (globalThis as { fetch: typeof fetch }).fetch = fetchMock as unknown as typeof fetch; + + const response = await fetchWithRetry( + "https://example.com", + { method: "POST", body: JSON.stringify({ hello: "world" }) }, + "acct-1", + ); + + expect(response.status).toBe(429); + expect(response.headers.get("X-Gemini-Terminal-429")).toBe("true"); + expect(response.headers.get("X-Gemini-429-Reason")).toBe("QUOTA_EXHAUSTED"); + expect(fetchMock.mock.calls.length).toBe(1); + }); + + it("includes terminal 429 headers for model capacity exhaustion", async () => { + const fetchMock = mock(async () => makeQuota429WithMessage( + "MODEL_CAPACITY_EXHAUSTED", + "No capacity available for model gemini-3-flash-preview on the server", + true, + )); + (globalThis as { fetch: typeof fetch }).fetch = fetchMock as unknown as typeof fetch; + + const response = await fetchWithRetry( + "https://example.com", + { method: "POST", body: JSON.stringify({ hello: "world" }) }, + "acct-1", + ); + + expect(response.status).toBe(429); + expect(response.headers.get("X-Gemini-Terminal-429")).toBe("true"); + expect(response.headers.get("X-Gemini-429-Reason")).toBe("MODEL_CAPACITY_EXHAUSTED"); + }); + + it("passes accountId through to retry throttle key", () => { + const keyWith = buildRetryThrottleKey( + "https://example.com", + { body: JSON.stringify({ project: "p1", model: "m1" }) }, + "acct-1", + ); + const keyWithout = buildRetryThrottleKey( + "https://example.com", + { body: JSON.stringify({ project: "p1", model: "m1" }) }, + ); + + expect(keyWith).toContain("account:acct-1"); + expect(keyWithout).not.toContain("account:"); + expect(keyWith).not.toBe(keyWithout); + }); +}); + +describe("formatAllAccountsExhaustedMessage", () => { + it("formats single account with no cooldown", () => { + const now = Date.now(); + const accounts: AccountState[] = [{ + account: { id: "user1", refreshToken: "rt1", enabled: true }, + auth: { type: "oauth", refresh: "rt1" }, + cooldowns: new Map(), + lastUsed: now, + usageCount: 0, + health: { value: 1.0, successRate: 1.0, quotaRemaining: 1.0, latencyScore: 1.0, cooldownScore: 1.0 }, + }]; + + const result = formatAllAccountsExhaustedMessage(accounts); + expect(result).toContain("All accounts exhausted."); + expect(result).toContain("user1"); + expect(result).toContain("no cooldown"); + }); + + it("lists active cooldowns with remaining time", () => { + const now = Date.now(); + const cooldown: CooldownState = { + accountId: "user1", + expiresAt: now + 30000, + reason: "QUOTA_EXHAUSTED", + model: "gemini-2.5-flash", + }; + const accounts: AccountState[] = [{ + account: { id: "user1", refreshToken: "rt1", enabled: true }, + auth: { type: "oauth", refresh: "rt1" }, + cooldowns: new Map([["gemini-2.5-flash", cooldown]]), + lastUsed: now, + usageCount: 0, + health: { value: 1.0, successRate: 1.0, quotaRemaining: 1.0, latencyScore: 1.0, cooldownScore: 1.0 }, + }]; + + const result = formatAllAccountsExhaustedMessage(accounts); + expect(result).toContain("user1"); + expect(result).toContain("gemini-2.5-flash: QUOTA_EXHAUSTED"); + expect(result).toContain("30s"); + }); + + it("handles multiple accounts with mix of cooldowns and non-cooldowned", () => { + const now = Date.now(); + const accounts: AccountState[] = [ + { + account: { id: "user1", refreshToken: "rt1", enabled: true }, + auth: { type: "oauth", refresh: "rt1" }, + cooldowns: new Map(), + lastUsed: now, + usageCount: 0, + health: { value: 1.0, successRate: 1.0, quotaRemaining: 1.0, latencyScore: 1.0, cooldownScore: 1.0 }, + }, + { + account: { id: "user2", refreshToken: "rt2", enabled: true, email: "user2@test.com" }, + auth: { type: "oauth", refresh: "rt2" }, + cooldowns: new Map([ + ["model-x", { + accountId: "user2", + expiresAt: now + 10000, + reason: "MODEL_CAPACITY_EXHAUSTED", + model: "model-x", + } as CooldownState], + ]), + lastUsed: now, + usageCount: 0, + health: { value: 0.5, successRate: 0.5, quotaRemaining: 0.5, latencyScore: 1.0, cooldownScore: 0.5 }, + }, + ]; + + const result = formatAllAccountsExhaustedMessage(accounts); + expect(result).toContain("user1: no cooldown"); + expect(result).toContain("user2"); + expect(result).toContain("model-x: MODEL_CAPACITY_EXHAUSTED"); + expect(result).toContain("10s"); + expect(result).toContain("All accounts exhausted."); + }); + + it("excludes expired cooldowns from output", () => { + const now = Date.now(); + const expiredCooldown: CooldownState = { + accountId: "user1", + expiresAt: now - 1000, + reason: "QUOTA_EXHAUSTED", + model: "gemini-2.5-flash", + }; + const accounts: AccountState[] = [{ + account: { id: "user1", refreshToken: "rt1", enabled: true }, + auth: { type: "oauth", refresh: "rt1" }, + cooldowns: new Map([["gemini-2.5-flash", expiredCooldown]]), + lastUsed: now, + usageCount: 0, + health: { value: 1.0, successRate: 1.0, quotaRemaining: 1.0, latencyScore: 1.0, cooldownScore: 1.0 }, + }]; + + const result = formatAllAccountsExhaustedMessage(accounts); + expect(result).toContain("no cooldown"); + expect(result).not.toContain("QUOTA_EXHAUSTED"); + }); }); describe("retryInternals", () => { diff --git a/src/plugin/retry/index.ts b/src/plugin/retry/index.ts index 856d80e..24bc678 100644 --- a/src/plugin/retry/index.ts +++ b/src/plugin/retry/index.ts @@ -10,6 +10,7 @@ import { import { classifyQuotaResponse, retryInternals } from "./quota"; import { isGeminiDebugEnabled, logGeminiDebugMessage } from "../debug"; import { geminiFetch } from "../../fetch"; +import type { AccountState } from "../types"; const retryCooldownByKey = new Map(); const RETRY_IN_FLIGHT_LOG_INTERVAL_MS = 5000; @@ -24,13 +25,14 @@ const MODEL_CAPACITY_COOLDOWN_MS = 8000; export async function fetchWithRetry( input: RequestInfo, init: RequestInit | undefined, + accountId?: string, ): Promise { if (!canRetryRequest(init)) { return geminiFetch(input, init); } const retryInit = cloneRetryableInit(init); - const throttleKey = buildRetryThrottleKey(input, retryInit); + const throttleKey = buildRetryThrottleKey(input, retryInit, accountId); await waitForRetryCooldown(throttleKey, retryInit.signal); let attempt = 1; const url = readRequestUrl(input); @@ -66,13 +68,15 @@ export async function fetchWithRetry( } stopInFlightLog(); - if (!isRetryableStatus(response.status)) { + if (!isRetryableStatus(response.status) && response.status !== 400 && response.status !== 403) { debugRetry(`attempt ${attempt} success or non-retryable status: ${response.status}`); return response; } - const quotaContext = response.status === 429 ? await classifyQuotaResponse(response) : null; - if (response.status === 429 && quotaContext?.terminal) { + const isPotentialQuotaError = response.status === 429 || response.status === 400 || response.status === 403; + const quotaContext = isPotentialQuotaError ? await classifyQuotaResponse(response) : null; + + if (isPotentialQuotaError && quotaContext?.terminal) { if (quotaContext.reason === "MODEL_CAPACITY_EXHAUSTED") { const cooldownMs = quotaContext.retryDelayMs ?? MODEL_CAPACITY_COOLDOWN_MS; setRetryCooldown(throttleKey, cooldownMs); @@ -81,6 +85,17 @@ export async function fetchWithRetry( debugRetry( `attempt ${attempt} terminal 429 (${quotaContext.reason ?? "unknown"}), returning without retry`, ); + // Add headers so callers (e.g. plugin.ts with AccountPool) can cooldown accounts + if (accountId) { + const headers = new Headers(response.headers); + headers.set("X-Gemini-Terminal-429", "true"); + headers.set("X-Gemini-429-Reason", quotaContext.reason ?? "unknown"); + return new Response(response.body, { + status: response.status, + statusText: response.statusText, + headers, + }); + } return response; } @@ -117,12 +132,13 @@ function cloneRetryableInit(init: RequestInit | undefined): RequestInit { }; } -function buildRetryThrottleKey(input: RequestInfo, init: RequestInit): string { +export function buildRetryThrottleKey(input: RequestInfo, init: RequestInit, accountId?: string): string { const url = readRequestUrl(input); const body = typeof init.body === "string" ? safeParseBody(init.body) : null; const project = readString(body?.project); const model = readString(body?.model); - return `${url}|${project ?? ""}|${model ?? ""}`; + const accountPrefix = accountId ? `account:${accountId}|` : ""; + return `${accountPrefix}${url}|${project ?? ""}|${model ?? ""}`; } async function waitForRetryCooldown(key: string, signal?: AbortSignal | null): Promise { @@ -220,4 +236,20 @@ function startInFlightLog(attempt: number, url: string): () => void { }; } +/** + * Formats a human-readable message when all pool accounts have exhausted their quota + * or are in cooldown. Used by plugin.ts when pool.select() returns an account that's + * already cooldowned and all others are too. + */ +export function formatAllAccountsExhaustedMessage(accounts: AccountState[]): string { + const now = Date.now(); + const lines = accounts.map((a) => { + const cooldowns = Array.from(a.cooldowns.values()) + .filter((c) => c.expiresAt > now) + .map((c) => `${c.model ?? "all"}: ${c.reason} (${Math.round((c.expiresAt - now) / 1000)}s)`); + return ` ${a.account.id}: ${cooldowns.join(", ") || "no cooldown"}`; + }); + return `All accounts exhausted. Cooldown status:\n${lines.join("\n")}`; +} + export { retryInternals }; diff --git a/src/plugin/retry/quota.ts b/src/plugin/retry/quota.ts index 1890969..93c6b16 100644 --- a/src/plugin/retry/quota.ts +++ b/src/plugin/retry/quota.ts @@ -100,6 +100,13 @@ export async function classifyQuotaResponse(response: Response): Promise = {}): HealthScore { + return { + value: 1.0, + successRate: 1.0, + quotaRemaining: 1.0, + latencyScore: 1.0, + cooldownScore: 1.0, + ...overrides, + }; +} + +function makeCooldown( + accountId: string, + expiresAt: number, + reason = "RATE_LIMIT", + model?: string, +): CooldownState { + return { accountId, expiresAt, reason, model }; +} + +function makeAccountState( + id: string, + overrides: Partial = {}, +): AccountState { + return { + account: { + id, + email: `${id}@example.com`, + refreshToken: `rt-${id}`, + enabled: true, + }, + auth: { type: "oauth", refresh: `rt-${id}`, access: "tok", expires: BASE_TIME + 3600000 }, + cooldowns: new Map(), + lastUsed: 0, + usageCount: 0, + health: makeScore(), + ...overrides, + } satisfies AccountState; +} + +describe("selectBestAccount", () => { + it("returns undefined for an empty pool", () => { + expect(selectBestAccount([])).toBeUndefined(); + }); + + it("returns the only account when the pool has one entry", () => { + const acc = makeAccountState("alice"); + expect(selectBestAccount([acc])).toBe(acc); + }); + + it("picks available account when others have zero health", () => { + const zero = makeAccountState("zero", { health: makeScore({ value: 0.0 }) }); + const high = makeAccountState("high", { health: makeScore({ value: 0.9 }) }); + const result = selectBestAccount([zero, high]); + expect(result?.account.id).toBe("high"); + }); + + it("canonicalizes model names in selectBestAccount", () => { + const now = Date.now(); + const acc = makeAccountState("acc", { + cooldowns: new Map([["models/gemini-1.5-pro", makeCooldown("acc", now + 60000, "RATE_LIMIT", "models/gemini-1.5-pro")]]) + }); + const other = makeAccountState("other"); + // gemini-1.5-pro should map to models/gemini-1.5-pro + const result = selectBestAccount([acc, other], { model: "gemini-1.5-pro" }); + expect(result?.account.id).toBe("other"); + }); + + it("picks the account with shortest remaining cooldown when all are cooldowned", () => { + const now = Date.now(); + const model = "models/gemini-2.0-flash"; + // long cooldown: expires far in future + const longCd = makeAccountState("long", { + cooldowns: new Map([[model, makeCooldown("long", now + 60000, "MODEL_CAPACITY_EXHAUSTED", model)]]), + }); + // short cooldown: expires sooner + const shortCd = makeAccountState("short", { + cooldowns: new Map([[model, makeCooldown("short", now + 10000, "MODEL_CAPACITY_EXHAUSTED", model)]]), + }); + const result = selectBestAccount([longCd, shortCd], { model: "gemini-2.0-flash" }); + expect(result?.account.id).toBe("short"); + }); + + it("never selects a disabled account", () => { + const disabled = makeAccountState("bob", { account: { id: "bob", email: "bob@x.com", refreshToken: "rt-bob", enabled: false } }); + const enabled = makeAccountState("alice"); + const result = selectBestAccount([disabled, enabled]); + expect(result?.account.id).toBe("alice"); + }); + + it("returns undefined when all accounts are disabled", () => { + const a = makeAccountState("a", { account: { id: "a", email: "a@x.com", refreshToken: "rt-a", enabled: false } }); + const b = makeAccountState("b", { account: { id: "b", email: "b@x.com", refreshToken: "rt-b", enabled: false } }); + expect(selectBestAccount([a, b])).toBeUndefined(); + }); + + it("prefers available healthy account over cooldowned ones", () => { + const now = Date.now(); + const model = "models/gemini-2.0-flash"; + const cooldowned = makeAccountState("coold", { + cooldowns: new Map([[model, makeCooldown("coold", now + 30000, "RATE_LIMIT", model)]]), + }); + const healthy = makeAccountState("healthy", { health: makeScore({ value: 0.95 }) }); + const result = selectBestAccount([cooldowned, healthy], { model: "gemini-2.0-flash" }); + expect(result?.account.id).toBe("healthy"); + }); + + it("breaks ties by LRU when weights are zero", () => { + const busy = makeAccountState("busy", { usageCount: 10, health: makeScore({ value: 0.0 }) }); + const idle = makeAccountState("idle", { usageCount: 2, health: makeScore({ value: 0.0 }) }); + const result = selectBestAccount([busy, idle]); + expect(result?.account.id).toBe("idle"); + }); + + it("prefers accounts with >20% quota remaining over those with low quota", () => { + const lowQuota = makeAccountState("low-quota", { health: makeScore({ value: 0.8, quotaRemaining: 0.1 }) }); + const highQuota = makeAccountState("high-quota", { health: makeScore({ value: 0.7, quotaRemaining: 0.5 }) }); + const result = selectBestAccount([lowQuota, highQuota]); + // high-quota has lower health but >20% quota; low-quota has higher health but <20% + // Quota filter prefers >20% first; within those, highest health + expect(result?.account.id).toBe("high-quota"); + }); + + it("uses weighted random when all accounts have <20% quota", () => { + const zero = makeAccountState("zero-health", { health: makeScore({ value: 0.0, quotaRemaining: 0.1 }) }); + const higher = makeAccountState("higher-health", { health: makeScore({ value: 0.6, quotaRemaining: 0.05 }) }); + const result = selectBestAccount([zero, higher]); + // zero-health has 0 weight, so higher-health must be picked + expect(result?.account.id).toBe("higher-health"); + }); + + it("honors global 'all' cooldown even if specific model is requested", () => { + const now = Date.now(); + const globalCooldowned = makeAccountState("global-cd", { + cooldowns: new Map([["all", makeCooldown("global-cd", now + 60000, "AUTH_ERROR", "all")]]) + }); + const healthy = makeAccountState("healthy"); + + // Requested model gemini-1.5-pro should be blocked by 'all' cooldown + const result = selectBestAccount([globalCooldowned, healthy], { model: "gemini-1.5-pro" }); + expect(result?.account.id).toBe("healthy"); + }); +}); + +describe("explainSelection", () => { + it("returns a readable summary for a normal selection", () => { + const a = makeAccountState("alice", { health: makeScore({ value: 0.9 }) }); + const b = makeAccountState("bob", { health: makeScore({ value: 0.5 }) }); + const selected = selectBestAccount([a, b]); + const summary = explainSelection([a, b], selected); + expect(summary).toContain("selected: "); + expect(summary).toContain(selected!.account.id); + }); + + it("reports no accounts available when selected is undefined", () => { + const summary = explainSelection([], undefined); + expect(summary).toBe("No accounts available"); + }); + + it("reports disabled accounts", () => { + const disabled = makeAccountState("disabled", { + account: { id: "disabled", email: "d@x.com", refreshToken: "rt-d", enabled: false }, + }); + const healthy = makeAccountState("healthy", { health: makeScore({ value: 0.9 }) }); + const selected = selectBestAccount([disabled, healthy]); + const summary = explainSelection([disabled, healthy], selected, {}); + expect(summary).toContain("1 disabled"); + expect(summary).toContain("selected: healthy"); + }); + + it("reports global cooldowns in explainSelection", () => { + const now = Date.now(); + const globalCd = makeAccountState("global-cd", { + cooldowns: new Map([["all", makeCooldown("global-cd", now + 60000, "AUTH_ERROR", "all")]]) + }); + const healthy = makeAccountState("healthy"); + const summary = explainSelection([globalCd, healthy], healthy, { model: "gemini-1.5-pro" }); + expect(summary).toContain("1 in cooldown"); + expect(summary).toContain("selected: healthy"); + }); +}); diff --git a/src/plugin/rotation.ts b/src/plugin/rotation.ts new file mode 100644 index 0000000..93f7ad2 --- /dev/null +++ b/src/plugin/rotation.ts @@ -0,0 +1,111 @@ +/** + * Selects the best account from a pool using: + * 1. Filter out disabled accounts + * 2. Filter out accounts in cooldown for the given model + * 3. If all filtered out → pick account with shortest remaining cooldown + * 4. Apply health score weighting (higher health = higher probability) + * 5. Prefer accounts with >20% remaining quota (if quota data is fresh) + * 6. LRU tiebreaker (lowest usageCount) + */ + +import type { AccountState, CooldownState } from "./types"; +import { getCanonicalModelName } from "./models"; + +export interface SelectOptions { + model?: string; + url?: string; +} + +export function selectBestAccount( + accounts: AccountState[], + options: SelectOptions = {}, +): AccountState | undefined { + if (accounts.length === 0) return undefined; + if (accounts.length === 1) return accounts[0]; + + const enabled = accounts.filter((a) => a.account.enabled); + if (enabled.length === 0) return undefined; + + const now = Date.now(); + const model = options.model ? getCanonicalModelName(options.model) : undefined; + + // Check cooldown status per account + const cooldownStatus = enabled.map((account) => { + const modelCooldown = model ? account.cooldowns.get(model) : undefined; + const globalCooldown = account.cooldowns.get("all"); + + const isModelCooldowned = modelCooldown && modelCooldown.expiresAt > now; + const isGlobalCooldowned = globalCooldown && globalCooldown.expiresAt > now; + + const isCooldowned = isModelCooldowned || isGlobalCooldowned; + const remainingMs = Math.max( + isModelCooldowned ? modelCooldown!.expiresAt - now : 0, + isGlobalCooldowned ? globalCooldown!.expiresAt - now : 0, + ); + + return { account, isCooldowned: !!isCooldowned, remainingMs }; + }); + + const available = cooldownStatus.filter((c) => !c.isCooldowned); + + // If all are cooldowned, pick the one with shortest remaining cooldown + if (available.length === 0) { + const shortest = cooldownStatus.reduce((a, b) => + a.remainingMs < b.remainingMs ? a : b, + ); + return shortest.account; + } + + // Filter by quota preference (>20% remaining) + const quotaPreferred = available.filter( + (c) => c.account.health.quotaRemaining > 0.2, + ); + const candidates = quotaPreferred.length > 0 ? quotaPreferred : available; + + // Weighted random selection among accounts with health > 0 + const totalHealth = candidates.reduce((sum, c) => sum + Math.max(0, c.account.health.value), 0); + + if (totalHealth > 0) { + let random = Math.random() * totalHealth; + for (const c of candidates) { + const weight = Math.max(0, c.account.health.value); + if (random <= weight) { + return c.account; + } + random -= weight; + } + } + + // Fallback: pick the one with lowest usageCount (LRU) + return candidates.reduce((a, b) => + a.account.usageCount <= b.account.usageCount ? a : b, + ).account; +} + +/** + * Returns a human-readable summary of why an account was or wasn't selected. + */ +export function explainSelection( + accounts: AccountState[], + selected: AccountState | undefined, + options: SelectOptions = {}, +): string { + if (!selected) return "No accounts available"; + if (accounts.length === 1) return `Single account: ${selected.account.id}`; + + const now = Date.now(); + const model = options.model; + const disabled = accounts.filter((a) => !a.account.enabled); + const cooldowned = accounts.filter((a) => { + const modelCd = model ? a.cooldowns.get(getCanonicalModelName(model)) : undefined; + const globalCd = a.cooldowns.get("all"); + return (modelCd && modelCd.expiresAt > now) || (globalCd && globalCd.expiresAt > now); + }); + + const parts: string[] = []; + if (disabled.length > 0) parts.push(`${disabled.length} disabled`); + if (cooldowned.length > 0) parts.push(`${cooldowned.length} in cooldown`); + parts.push(`selected: ${selected.account.id} (health: ${selected.health.value})`); + + return parts.join(", "); +} diff --git a/src/plugin/token.test.ts b/src/plugin/token.test.ts index c991cbf..4cc893f 100644 --- a/src/plugin/token.test.ts +++ b/src/plugin/token.test.ts @@ -1,8 +1,8 @@ import { afterEach, beforeEach, describe, expect, it, mock } from "bun:test"; import { GEMINI_PROVIDER_ID } from "../constants"; -import { refreshAccessToken } from "./token"; -import type { OAuthAuthDetails, PluginClient } from "./types"; +import { refreshAccessToken, refreshAccessTokenForAccount } from "./token"; +import type { AccountPool, OAuthAuthDetails, PluginClient } from "./types"; const originalSetTimeout = globalThis.setTimeout; @@ -159,3 +159,132 @@ describe("refreshAccessToken", () => { expect(fetchMock.mock.calls.length).toBe(2); }); }); + +describe("refreshAccessTokenForAccount", () => { + beforeEach(() => { + mock.restore(); + (globalThis as { setTimeout: typeof setTimeout }).setTimeout = ((fn: (...args: any[]) => void) => { + fn(); + return 0 as unknown as ReturnType; + }) as typeof setTimeout; + }); + + afterEach(() => { + (globalThis as { setTimeout: typeof setTimeout }).setTimeout = originalSetTimeout; + }); + + it("calls pool.withRefreshLock with the correct accountId", async () => { + const withRefreshLock = mock(async (_accountId: string, fn: () => Promise) => fn()); + const pool = { + getAccount: mock(() => ({ + account: { id: "acc@test.com" }, + auth: { type: "oauth" as const, refresh: "refresh-token" }, + })), + withRefreshLock, + updateAuth: mock(() => {}), + getAccounts: mock(() => [{ account: { id: "acc@test.com" } }]), + } as unknown as AccountPool; + + const client = createClient(); + const fetchMock = mock(async () => { + return new Response( + JSON.stringify({ access_token: "new-access", expires_in: 3600 }), + { status: 200 }, + ); + }); + (globalThis as { fetch: typeof fetch }).fetch = fetchMock as unknown as typeof fetch; + + await refreshAccessTokenForAccount(pool, "acc@test.com", client); + + expect(withRefreshLock.mock.calls.length).toBe(1); + expect(withRefreshLock.mock.calls[0]?.[0]).toBe("acc@test.com"); + }); + + it("updates pool state on successful refresh", async () => { + const updateAuth = mock(() => {}); + const pool = { + getAccount: mock(() => ({ + account: { id: "acc@test.com" }, + auth: { type: "oauth" as const, refresh: "refresh-token" }, + })), + withRefreshLock: mock(async (_accountId: string, fn: () => Promise) => fn()), + updateAuth, + getAccounts: mock(() => [{ account: { id: "acc@test.com" } }]), + } as unknown as AccountPool; + + const client = createClient(); + const fetchMock = mock(async () => { + return new Response( + JSON.stringify({ access_token: "new-access", expires_in: 3600 }), + { status: 200 }, + ); + }); + (globalThis as { fetch: typeof fetch }).fetch = fetchMock as unknown as typeof fetch; + + const result = await refreshAccessTokenForAccount(pool, "acc@test.com", client); + + expect(result?.access).toBe("new-access"); + expect(updateAuth.mock.calls.length).toBe(1); + expect(updateAuth.mock.calls[0]?.[0]).toBe("acc@test.com"); + }); + + it("returns null for non-existent account", async () => { + const pool = { + getAccount: mock(() => undefined), + withRefreshLock: mock(() => {}), + updateAuth: mock(() => {}), + getAccounts: mock(() => [{ account: { id: "acc@test.com" } }]), + } as unknown as AccountPool; + + const client = createClient(); + const result = await refreshAccessTokenForAccount(pool, "missing@test.com", client); + + expect(result).toBeNull(); + }); + + it("deduplicates concurrent refresh calls via pool lock", async () => { + let release!: () => void; + const gate = new Promise((resolve) => { release = resolve; }); + + // Simulate pool-level dedup: cache the promise returned by withRefreshLock + let lockPromise: Promise | null = null; + const withRefreshLock = mock(async (_accountId: string, fn: () => Promise) => { + if (!lockPromise) { + lockPromise = fn(); + } + return lockPromise; + }); + + const pool = { + getAccount: mock(() => ({ + account: { id: "acc@test.com" }, + auth: { type: "oauth" as const, refresh: "refresh-token" }, + })), + withRefreshLock, + updateAuth: mock(() => {}), + getAccounts: mock(() => [{ account: { id: "acc@test.com" } }]), + } as unknown as AccountPool; + + const client = createClient(); + const fetchMock = mock(async () => { + await gate; + return new Response( + JSON.stringify({ access_token: "deduped-access", expires_in: 3600 }), + { status: 200 }, + ); + }); + (globalThis as { fetch: typeof fetch }).fetch = fetchMock as unknown as typeof fetch; + + const first = refreshAccessTokenForAccount(pool, "acc@test.com", client); + const second = refreshAccessTokenForAccount(pool, "acc@test.com", client); + await Promise.resolve(); + + // fetch should be called only once (second call returned cached lock promise) + expect(fetchMock.mock.calls.length).toBe(1); + release!(); + + const [firstResult, secondResult] = await Promise.all([first, second]); + expect(firstResult?.access).toBe("deduped-access"); + expect(secondResult?.access).toBe("deduped-access"); + }); +}); diff --git a/src/plugin/token.ts b/src/plugin/token.ts index 5fa7600..4de7ee8 100644 --- a/src/plugin/token.ts +++ b/src/plugin/token.ts @@ -20,18 +20,8 @@ import { resolveRetryDelayMs, wait, } from "./retry/helpers"; -import type { OAuthAuthDetails, PluginClient, RefreshParts } from "./types"; - -interface OAuthErrorPayload { - error?: - | string - | { - code?: string; - status?: string; - message?: string; - }; - error_description?: string; -} +import { storeCachedAuthForAccount } from "./cache"; +import { saveAccountsToDisk } from "./config-store"; const refreshInFlight = new Map>(); @@ -249,3 +239,34 @@ async function fetchTokenRefresh(refreshToken: string): Promise { return geminiFetch(tokenUrl, init); } + +/** + * Refreshes the access token for a specific account, updating the pool state directly. + * Uses the pool's withRefreshLock to prevent concurrent refresh races. + */ +export async function refreshAccessTokenForAccount( + pool: any, + accountId: string, + client: PluginClient, +): Promise { + const state = pool.getAccount(accountId); + if (!state) return null; + + return pool.withRefreshLock(accountId, async () => { + const currentAuth = state.auth; + if (!currentAuth.refresh) return null; + + const refreshed = await refreshAccessToken(currentAuth, client); + if (refreshed) { + pool.updateAuth(accountId, refreshed); + storeCachedAuthForAccount(accountId, refreshed); + + // Persist the updated refresh token to disk + const accountsToSave = pool.getAccounts().map((s: any) => s.account); + await saveAccountsToDisk(accountsToSave).catch(e => { + console.error("[Gemini Auth] Failed to persist refreshed token:", e); + }); + } + return refreshed ?? null; + }); +} diff --git a/src/plugin/types.ts b/src/plugin/types.ts index a0ac506..2db5ff3 100644 --- a/src/plugin/types.ts +++ b/src/plugin/types.ts @@ -97,3 +97,46 @@ export interface ProjectContextResult { auth: OAuthAuthDetails; effectiveProjectId: string; } + +// ── Multi-account types ────────────────────────────────────────────── + +export interface GeminiAccount { + id: string; // email from OAuth userinfo + email?: string; + refreshToken: string; + projectId?: string; // optional per-account project override + enabled: boolean; +} + +export interface HealthScore { + value: number; // 0.0 - 1.0 + successRate: number; // 0.0 - 1.0 + quotaRemaining: number; // 0.0 - 1.0 + latencyScore: number; // 0.0 - 1.0 + cooldownScore: number; // 0.0 - 1.0 +} + +export interface CooldownState { + accountId: string; + expiresAt: number; // Date.now() + durationMs + reason: string; // e.g. "MODEL_CAPACITY_EXHAUSTED", "QUOTA_EXHAUSTED" + model?: string; +} + +export type RotationStrategy = "round-robin" | "lru" | "quota-aware" | "health-weighted"; + +export interface AccountPoolConfig { + accounts: GeminiAccount[]; + strategy?: RotationStrategy; +} + +export interface AccountState { + account: GeminiAccount; + auth: OAuthAuthDetails; + projectContext?: ProjectContextResult; + cooldowns: Map; // key: model or "model|url" + lastUsed: number; // timestamp + usageCount: number; + health: HealthScore; + refreshLock?: Promise; // per-account refresh lock +}