From b9f4f87c90b15285712a2d5b74f7f786aa6e68e8 Mon Sep 17 00:00:00 2001 From: Deuri Vasquez Date: Thu, 21 May 2026 06:39:53 -0400 Subject: [PATCH 1/9] feat: add Used column to Gemini quota output Shows consumed quota per model when remainingFraction and remainingAmount are both available, calculated as total = remaining / fraction, used = total - remaining. --- src/plugin/quota.test.ts | 5 ++++ src/plugin/quota.ts | 49 ++++++++++++++++++++++++++++------------ 2 files changed, 39 insertions(+), 15 deletions(-) diff --git a/src/plugin/quota.test.ts b/src/plugin/quota.test.ts index 13b17d5..5e0c09b 100644 --- a/src/plugin/quota.test.ts +++ b/src/plugin/quota.test.ts @@ -69,6 +69,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 +83,9 @@ 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) expect(output.indexOf("Gemini 3 (1 model, 1 bucket)")).toBeLessThan( output.indexOf("Gemini 2.5 (2 models, 3 buckets)"), ); @@ -113,6 +117,7 @@ describe("formatGeminiQuotaOutput", () => { const output = formatGeminiQuotaOutput("test-project", buckets); expect(output).toContain("Type"); + expect(output).toContain("Used"); expect(output).toContain("Gemini 2.5 (1 model, 2 buckets)"); expect(output).toContain("REQUESTS"); expect(output).toContain("TOKENS"); diff --git a/src/plugin/quota.ts b/src/plugin/quota.ts index 1333eeb..fd273ce 100644 --- a/src/plugin/quota.ts +++ b/src/plugin/quota.ts @@ -102,8 +102,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 +120,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 +149,42 @@ 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}%`; + return { + used: "unknown", + 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 +259,7 @@ function normalizeTokenType(bucket: RetrieveUserQuotaBucket): string { interface GroupedQuotaRow { variant: string; + usageUsed: string; usageRemaining: string; resetValue: string; tokenType: string; @@ -260,7 +277,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 +286,8 @@ function groupQuotaRows(sortedBuckets: RetrieveUserQuotaBucket[]): GroupedQuotaM if (existing) { existing.rows.push({ variant, - usageRemaining, + usageUsed: usageInfo.used, + usageRemaining: usageInfo.remaining, resetValue, tokenType, }); @@ -281,7 +299,8 @@ function groupQuotaRows(sortedBuckets: RetrieveUserQuotaBucket[]): GroupedQuotaM version: extractModelVersion(baseModel), rows: [{ variant, - usageRemaining, + usageUsed: usageInfo.used, + usageRemaining: usageInfo.remaining, resetValue, tokenType, }], From 285fd32cd95648e861738fe2ce54f29a2e9c5bc0 Mon Sep 17 00:00:00 2001 From: Deuri Vasquez Date: Thu, 21 May 2026 06:56:22 -0400 Subject: [PATCH 2/9] fix: show used percentage when API returns only fraction The Code Assist quota API returns remainingFraction but not remainingAmount, so absolute 'used' counts cannot be calculated. Fall back to showing used percentage (100% - remaining%) instead of 'unknown' when only fraction data is available. --- src/plugin/quota.test.ts | 5 +++++ src/plugin/quota.ts | 3 ++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/src/plugin/quota.test.ts b/src/plugin/quota.test.ts index 5e0c09b..0cbf2cd 100644 --- a/src/plugin/quota.test.ts +++ b/src/plugin/quota.test.ts @@ -86,6 +86,9 @@ describe("formatGeminiQuotaOutput", () => { // 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)"), ); @@ -118,6 +121,8 @@ 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"); diff --git a/src/plugin/quota.ts b/src/plugin/quota.ts index fd273ce..03d4710 100644 --- a/src/plugin/quota.ts +++ b/src/plugin/quota.ts @@ -174,8 +174,9 @@ function formatUsageInfo(bucket: RetrieveUserQuotaBucket): UsageInfo { if (hasFraction) { const clamped = clamp(remainingFraction, 0, 1); + const usedPercent = ((1 - clamped) * 100).toFixed(1); return { - used: "unknown", + used: `${usedPercent}%`, remaining: `${buildProgressBar(clamped)} ${(clamped * 100).toFixed(1)}%`, }; } From 3242e7de1f16463ebe9e33be613f3ea46b9507f7 Mon Sep 17 00:00:00 2001 From: Deuri Vasquez Date: Thu, 21 May 2026 08:16:39 -0400 Subject: [PATCH 3/9] feat: add multi-account support with automatic model rotation - Introduce AccountPool for managing multiple accounts - Implement health-weighted rotation strategy with quota fallback - Add AccountManager for dynamic account addition/removal - Persist dynamic accounts to ~/.config/opencode/gemini-auth.json - Make retry cooldowns, auth caches, and project contexts per-account - Update quota tool to display per-account and aggregate summaries --- README.md | 57 ++++ .../changes/multi-account-rotation/tasks.md | 52 ++++ openspec/config.yaml | 33 +++ src/plugin.ts | 260 ++++++++++++++---- src/plugin/account-manager.test.ts | 120 ++++++++ src/plugin/account-manager.ts | 122 ++++++++ src/plugin/account-pool.test.ts | 213 ++++++++++++++ src/plugin/account-pool.ts | 170 ++++++++++++ src/plugin/cache.ts | 16 ++ src/plugin/config-store.ts | 47 ++++ src/plugin/health.test.ts | 121 ++++++++ src/plugin/health.ts | 84 ++++++ src/plugin/oauth-authorize.ts | 77 +++++- src/plugin/plugin.test.ts | 128 +++++++++ src/plugin/project/context.test.ts | 176 ++++++++++++ src/plugin/project/context.ts | 28 ++ src/plugin/project/index.ts | 2 + src/plugin/project/utils.ts | 8 + src/plugin/provider.ts | 41 ++- src/plugin/quota.test.ts | 73 ++++- src/plugin/quota.ts | 127 ++++++++- src/plugin/retry.test.ts | 175 +++++++++++- src/plugin/retry/index.ts | 44 ++- src/plugin/retry/quota.ts | 7 + src/plugin/rotation.test.ts | 152 ++++++++++ src/plugin/rotation.ts | 102 +++++++ src/plugin/token.test.ts | 133 ++++++++- src/plugin/token.ts | 45 ++- src/plugin/types.ts | 43 +++ 29 files changed, 2569 insertions(+), 87 deletions(-) create mode 100644 openspec/changes/multi-account-rotation/tasks.md create mode 100644 openspec/config.yaml create mode 100644 src/plugin/account-manager.test.ts create mode 100644 src/plugin/account-manager.ts create mode 100644 src/plugin/account-pool.test.ts create mode 100644 src/plugin/account-pool.ts create mode 100644 src/plugin/config-store.ts create mode 100644 src/plugin/health.test.ts create mode 100644 src/plugin/health.ts create mode 100644 src/plugin/plugin.test.ts create mode 100644 src/plugin/project/context.test.ts create mode 100644 src/plugin/rotation.test.ts create mode 100644 src/plugin/rotation.ts 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..35f01e0 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,142 @@ 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)) { + const refreshed = await refreshAccessTokenForAccount(latestGeminiPool!, currentAccountId, client); + if (refreshed) { + selected = latestGeminiPool!.getAccount(currentAccountId); + } + 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 success/failure to pool + if (response.ok) { + latestGeminiPool!.reportSuccess(selected.account.id); + break; // Success! Exit the loop. + } else { + latestGeminiPool!.reportFailure(selected.account.id); + // 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"; + latestGeminiPool!.cooldownAccount(selected.account.id, model ?? "all", 8000, 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 +292,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 +448,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..536f279 --- /dev/null +++ b/src/plugin/account-pool.test.ts @@ -0,0 +1,213 @@ +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", () => { + const pool = createPool(2); + // Lower acc-0's health by recording a failure + 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); + // Both have equal health (1.0) → LRU tiebreaker picks acc-0 (first) + const before = pool.getAccount("acc-0")?.usageCount ?? -1; + const selected = pool.select(); + expect(selected?.account.id).toBe("acc-0"); + const after = pool.getAccount("acc-0")?.usageCount ?? -1; + expect(after).toBe(before + 1); + }); + }); + + describe("cooldown and health reporting", () => { + it("cooldownAccount sets a cooldown entry", () => { + 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("gemini-2.0-flash")).toBe(true); + expect(cooldowns?.get("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("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..124e584 --- /dev/null +++ b/src/plugin/account-pool.ts @@ -0,0 +1,170 @@ +/** + * 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 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; + state.cooldowns.set(model, { + accountId, + expiresAt: Date.now() + durationMs, + reason, + model, + }); + const tracker = this.trackers.get(accountId); + if (tracker) { + tracker.recordCooldown(); + state.health = tracker.compute(); + } + } + + 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 lock = refreshFn().finally(() => { + this.refreshLocks.delete(accountId); + }); + 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/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 0cbf2cd..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; @@ -129,3 +130,73 @@ describe("formatGeminiQuotaOutput", () => { 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 03d4710..95cfa81 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 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(), ); @@ -404,3 +417,111 @@ 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}`); + const accessToken = state.auth.access; + 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 the account with higher health score", () => { + const low = makeAccountState("low", { health: makeScore({ value: 0.4 }) }); + const high = makeAccountState("high", { health: makeScore({ value: 0.9 }) }); + const result = selectBestAccount([low, high]); + expect(result?.account.id).toBe("high"); + }); + + it("picks the account with shortest remaining cooldown when all are cooldowned", () => { + const now = Date.now(); + // long cooldown: expires far in future + const longCd = makeAccountState("long", { + cooldowns: new Map([["gemini-2.0-flash", makeCooldown("long", now + 60000, "MODEL_CAPACITY_EXHAUSTED", "gemini-2.0-flash")]]), + }); + // short cooldown: expires sooner + const shortCd = makeAccountState("short", { + cooldowns: new Map([["gemini-2.0-flash", makeCooldown("short", now + 10000, "MODEL_CAPACITY_EXHAUSTED", "gemini-2.0-flash")]]), + }); + 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 cooldowned = makeAccountState("coold", { + cooldowns: new Map([["gemini-2.0-flash", makeCooldown("coold", now + 30000, "RATE_LIMIT", "gemini-2.0-flash")]]), + }); + 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 (lowest usageCount wins)", () => { + const busy = makeAccountState("busy", { usageCount: 10, health: makeScore({ value: 0.9 }) }); + const idle = makeAccountState("idle", { usageCount: 2, health: makeScore({ value: 0.9 }) }); + 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 lowest health when all accounts have <20% quota since quota filter is skipped", () => { + const lower = makeAccountState("lower-health", { health: makeScore({ value: 0.3, quotaRemaining: 0.1 }) }); + const higher = makeAccountState("higher-health", { health: makeScore({ value: 0.6, quotaRemaining: 0.05 }) }); + const result = selectBestAccount([lower, higher]); + // both have <20% quota → skip quota filter → pick highest health + expect(result?.account.id).toBe("higher-health"); + }); +}); + +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: alice"); + expect(summary).toContain("health: 0.9"); + }); + + 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"); + }); +}); diff --git a/src/plugin/rotation.ts b/src/plugin/rotation.ts new file mode 100644 index 0000000..1d88362 --- /dev/null +++ b/src/plugin/rotation.ts @@ -0,0 +1,102 @@ +/** + * 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"; + +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; + + // Check cooldown status per account + const cooldownStatus = enabled.map((account) => { + const cooldown = model ? account.cooldowns.get(model) : undefined; + const isCooldowned = cooldown && cooldown.expiresAt > now; + const remainingMs = isCooldowned ? cooldown!.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; + + // Health-weighted selection: pick the account with highest health score + // For simplicity in this phase, use deterministic highest-health selection + // (probabilistic weighting can be added later) + const best = candidates.reduce((a, b) => + a.account.health.value >= b.account.health.value ? a : b, + ); + + // Among accounts with equal health, use LRU (lowest usageCount) + const topHealth = best.account.health.value; + const equalHealth = candidates.filter( + (c) => c.account.health.value === topHealth, + ); + if (equalHealth.length > 1) { + return equalHealth.reduce((a, b) => + a.account.usageCount <= b.account.usageCount ? a : b, + ).account; + } + + return best.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) => { + if (!model) return false; + const cd = a.cooldowns.get(model); + return cd && cd.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 +} From d33cdd3e4a0d3b9ccd6a39c44274ab5c43856b82 Mon Sep 17 00:00:00 2001 From: Deuri Vasquez Date: Thu, 21 May 2026 08:54:24 -0400 Subject: [PATCH 4/9] docs: add AGENTS.md with repository instructions --- AGENTS.md | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) create mode 100644 AGENTS.md 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. From 4f70b290662bfd71433c1c280424fb0211f3d9cd Mon Sep 17 00:00:00 2001 From: Deuri Vasquez Date: Thu, 21 May 2026 08:55:46 -0400 Subject: [PATCH 5/9] chore: ignore .atl directory --- .gitignore | 2 ++ 1 file changed, 2 insertions(+) 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 From d124fb68e78a2fb7ee365512d8f11177ed4cf249 Mon Sep 17 00:00:00 2001 From: Deuri Vasquez Date: Thu, 21 May 2026 20:38:24 -0400 Subject: [PATCH 6/9] feat(rotation): implement Judgment Day Round 1 fixes --- src/plugin.ts | 24 +++++++++++++--- src/plugin/account-pool.test.ts | 49 +++++++++++++++++++++++++-------- src/plugin/account-pool.ts | 26 ++++++++++++++--- src/plugin/models.ts | 39 ++++++++++++++++++++++++++ src/plugin/rotation.test.ts | 43 +++++++++++++++++++---------- src/plugin/rotation.ts | 36 ++++++++++++------------ 6 files changed, 165 insertions(+), 52 deletions(-) create mode 100644 src/plugin/models.ts diff --git a/src/plugin.ts b/src/plugin.ts index 35f01e0..6aedf64 100644 --- a/src/plugin.ts +++ b/src/plugin.ts @@ -236,16 +236,32 @@ export const GeminiCLIOAuthPlugin = async ( const response = await fetchWithRetry(transformed.request, transformed.init, selected.account.id); lastResponse = response; - // Report success/failure to pool + // Report result to pool (handles success, failure, and 401 circuit breaker) + latestGeminiPool!.reportResult(selected.account.id, response.status); + if (response.ok) { - latestGeminiPool!.reportSuccess(selected.account.id); break; // Success! Exit the loop. } else { - latestGeminiPool!.reportFailure(selected.account.id); // 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"; - latestGeminiPool!.cooldownAccount(selected.account.id, model ?? "all", 8000, reason); + + // 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); diff --git a/src/plugin/account-pool.test.ts b/src/plugin/account-pool.test.ts index 536f279..ea0df63 100644 --- a/src/plugin/account-pool.test.ts +++ b/src/plugin/account-pool.test.ts @@ -44,10 +44,15 @@ describe("AccountPool", () => { expect(selected?.account.id).toBe("acc-0"); }); - it("returns account with highest health", () => { + it("returns account with highest health (deterministic if other is zero)", () => { const pool = createPool(2); - // Lower acc-0's health by recording a failure - pool.reportFailure("acc-0"); + // 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"); @@ -78,22 +83,21 @@ describe("AccountPool", () => { it("increments usageCount and updates lastUsed on selection", () => { const pool = createPool(2); - // Both have equal health (1.0) → LRU tiebreaker picks acc-0 (first) - const before = pool.getAccount("acc-0")?.usageCount ?? -1; const selected = pool.select(); - expect(selected?.account.id).toBe("acc-0"); - const after = pool.getAccount("acc-0")?.usageCount ?? -1; - expect(after).toBe(before + 1); + 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", () => { + 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("gemini-2.0-flash")).toBe(true); - expect(cooldowns?.get("gemini-2.0-flash")?.reason).toBe("RATE_LIMIT"); + 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", () => { @@ -103,6 +107,29 @@ describe("AccountPool", () => { 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"); diff --git a/src/plugin/account-pool.ts b/src/plugin/account-pool.ts index 124e584..820e3b7 100644 --- a/src/plugin/account-pool.ts +++ b/src/plugin/account-pool.ts @@ -5,6 +5,7 @@ import { HealthTracker } from "./health"; import { selectBestAccount } from "./rotation"; +import { getCanonicalModelName } from "./models"; import type { AccountState, GeminiAccount, @@ -81,11 +82,12 @@ export class AccountPool { cooldownAccount(accountId: string, model: string, durationMs: number, reason: string): void { const state = this.states.get(accountId); if (!state) return; - state.cooldowns.set(model, { + const canonicalModel = getCanonicalModelName(model); + state.cooldowns.set(canonicalModel, { accountId, expiresAt: Date.now() + durationMs, reason, - model, + model: canonicalModel, }); const tracker = this.trackers.get(accountId); if (tracker) { @@ -94,6 +96,17 @@ export class AccountPool { } } + 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; @@ -152,9 +165,14 @@ export class AccountPool { const existing = this.refreshLocks.get(accountId); if (existing) return existing; - const lock = refreshFn().finally(() => { + 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; } 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/rotation.test.ts b/src/plugin/rotation.test.ts index a9131ff..8533647 100644 --- a/src/plugin/rotation.test.ts +++ b/src/plugin/rotation.test.ts @@ -55,22 +55,34 @@ describe("selectBestAccount", () => { expect(selectBestAccount([acc])).toBe(acc); }); - it("picks the account with higher health score", () => { - const low = makeAccountState("low", { health: makeScore({ value: 0.4 }) }); + 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([low, high]); + 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([["gemini-2.0-flash", makeCooldown("long", now + 60000, "MODEL_CAPACITY_EXHAUSTED", "gemini-2.0-flash")]]), + cooldowns: new Map([[model, makeCooldown("long", now + 60000, "MODEL_CAPACITY_EXHAUSTED", model)]]), }); // short cooldown: expires sooner const shortCd = makeAccountState("short", { - cooldowns: new Map([["gemini-2.0-flash", makeCooldown("short", now + 10000, "MODEL_CAPACITY_EXHAUSTED", "gemini-2.0-flash")]]), + 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"); @@ -91,17 +103,18 @@ describe("selectBestAccount", () => { 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([["gemini-2.0-flash", makeCooldown("coold", now + 30000, "RATE_LIMIT", "gemini-2.0-flash")]]), + 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 (lowest usageCount wins)", () => { - const busy = makeAccountState("busy", { usageCount: 10, health: makeScore({ value: 0.9 }) }); - const idle = makeAccountState("idle", { usageCount: 2, health: makeScore({ value: 0.9 }) }); + 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"); }); @@ -115,11 +128,11 @@ describe("selectBestAccount", () => { expect(result?.account.id).toBe("high-quota"); }); - it("uses lowest health when all accounts have <20% quota since quota filter is skipped", () => { - const lower = makeAccountState("lower-health", { health: makeScore({ value: 0.3, quotaRemaining: 0.1 }) }); + 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([lower, higher]); - // both have <20% quota → skip quota filter → pick highest health + const result = selectBestAccount([zero, higher]); + // zero-health has 0 weight, so higher-health must be picked expect(result?.account.id).toBe("higher-health"); }); }); @@ -130,8 +143,8 @@ describe("explainSelection", () => { const b = makeAccountState("bob", { health: makeScore({ value: 0.5 }) }); const selected = selectBestAccount([a, b]); const summary = explainSelection([a, b], selected); - expect(summary).toContain("selected: alice"); - expect(summary).toContain("health: 0.9"); + expect(summary).toContain("selected: "); + expect(summary).toContain(selected!.account.id); }); it("reports no accounts available when selected is undefined", () => { diff --git a/src/plugin/rotation.ts b/src/plugin/rotation.ts index 1d88362..400cd2d 100644 --- a/src/plugin/rotation.ts +++ b/src/plugin/rotation.ts @@ -9,6 +9,7 @@ */ import type { AccountState, CooldownState } from "./types"; +import { getCanonicalModelName } from "./models"; export interface SelectOptions { model?: string; @@ -26,7 +27,7 @@ export function selectBestAccount( if (enabled.length === 0) return undefined; const now = Date.now(); - const model = options.model; + const model = options.model ? getCanonicalModelName(options.model) : undefined; // Check cooldown status per account const cooldownStatus = enabled.map((account) => { @@ -52,25 +53,24 @@ export function selectBestAccount( ); const candidates = quotaPreferred.length > 0 ? quotaPreferred : available; - // Health-weighted selection: pick the account with highest health score - // For simplicity in this phase, use deterministic highest-health selection - // (probabilistic weighting can be added later) - const best = candidates.reduce((a, b) => - a.account.health.value >= b.account.health.value ? a : b, - ); - - // Among accounts with equal health, use LRU (lowest usageCount) - const topHealth = best.account.health.value; - const equalHealth = candidates.filter( - (c) => c.account.health.value === topHealth, - ); - if (equalHealth.length > 1) { - return equalHealth.reduce((a, b) => - a.account.usageCount <= b.account.usageCount ? a : b, - ).account; + // 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; + } } - return best.account; + // Fallback: pick the one with lowest usageCount (LRU) + return candidates.reduce((a, b) => + a.account.usageCount <= b.account.usageCount ? a : b, + ).account; } /** From dae0dffd91076e34e2544b9238bb3a947a3fb969 Mon Sep 17 00:00:00 2001 From: Deuri Vasquez Date: Thu, 21 May 2026 20:46:01 -0400 Subject: [PATCH 7/9] fix(rotation): handle refresh timeout gracefully --- src/plugin.ts | 24 +++++++++++++++++++++--- 1 file changed, 21 insertions(+), 3 deletions(-) diff --git a/src/plugin.ts b/src/plugin.ts index 6aedf64..6d39aff 100644 --- a/src/plugin.ts +++ b/src/plugin.ts @@ -169,10 +169,28 @@ export const GeminiCLIOAuthPlugin = async ( // Refresh token if needed if (accessTokenExpired(selected.auth)) { - const refreshed = await refreshAccessTokenForAccount(latestGeminiPool!, currentAccountId, client); - if (refreshed) { - selected = latestGeminiPool!.getAccount(currentAccountId); + 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"); From 6f0e4dc15df03c214b845d0b88b1b4264e974d96 Mon Sep 17 00:00:00 2001 From: Deuri Vasquez Date: Thu, 21 May 2026 20:46:01 -0400 Subject: [PATCH 8/9] fix(rotation): respect global 'all' cooldown in selection --- src/plugin/rotation.test.ts | 23 +++++++++++++++++++++++ src/plugin/rotation.ts | 21 +++++++++++++++------ 2 files changed, 38 insertions(+), 6 deletions(-) diff --git a/src/plugin/rotation.test.ts b/src/plugin/rotation.test.ts index 8533647..579cb94 100644 --- a/src/plugin/rotation.test.ts +++ b/src/plugin/rotation.test.ts @@ -135,6 +135,18 @@ describe("selectBestAccount", () => { // 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", () => { @@ -162,4 +174,15 @@ describe("explainSelection", () => { 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 index 400cd2d..93f7ad2 100644 --- a/src/plugin/rotation.ts +++ b/src/plugin/rotation.ts @@ -31,9 +31,18 @@ export function selectBestAccount( // Check cooldown status per account const cooldownStatus = enabled.map((account) => { - const cooldown = model ? account.cooldowns.get(model) : undefined; - const isCooldowned = cooldown && cooldown.expiresAt > now; - const remainingMs = isCooldowned ? cooldown!.expiresAt - now : 0; + 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 }; }); @@ -88,9 +97,9 @@ export function explainSelection( const model = options.model; const disabled = accounts.filter((a) => !a.account.enabled); const cooldowned = accounts.filter((a) => { - if (!model) return false; - const cd = a.cooldowns.get(model); - return cd && cd.expiresAt > now; + 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[] = []; From fc56fa8ee49866491f6c234f489231a441314a99 Mon Sep 17 00:00:00 2001 From: Deuri Vasquez Date: Thu, 21 May 2026 21:32:25 -0400 Subject: [PATCH 9/9] fix(quota): fetch access tokens for secondary accounts on demand --- src/plugin/quota.ts | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/src/plugin/quota.ts b/src/plugin/quota.ts index 95cfa81..f298d0f 100644 --- a/src/plugin/quota.ts +++ b/src/plugin/quota.ts @@ -3,7 +3,7 @@ import { accessTokenExpired, isOAuthAuth } from "./auth"; import { resolveCachedAuth } from "./cache"; import { ensureProjectContextForAccount, retrieveUserQuota } from "./project"; import type { RetrieveUserQuotaBucket } from "./project/types"; -import { refreshAccessToken } from "./token"; +import { refreshAccessToken, refreshAccessTokenForAccount } from "./token"; import type { AccountPool } from "./account-pool"; import type { AccountState, GetAuth, PluginClient } from "./types"; @@ -440,7 +440,20 @@ async function formatMultiAccountQuotaOutput( } lines.push(`Account: ${accountLabel}`); - const accessToken = state.auth.access; + 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("");