From 9f2ed1d58aeb17c9b4d39fb0cd3a6efd801ed38f Mon Sep 17 00:00:00 2001 From: roboomp Date: Tue, 9 Jun 2026 17:31:15 +0000 Subject: [PATCH] fix(ai): rotate antigravity credentials on Individual quota reached 429s - Extend USAGE_LIMIT_PATTERN with quota.?reached so auth-retry and AuthStorage.markUsageLimitReached recognise Antigravity's 'Individual quota reached' 429 as a credential-rotatable usage limit instead of a terminal provider error. The parseRateLimitReason classifier already mapped this phrasing to QUOTA_EXHAUSTED via the generic quota check. - Add antigravityRankingStrategy: picks the lowest-remainingFraction counter as primary and the next-lowest as secondary, with 24h windowDefaults matching the daily-cloudcode-pa.googleapis.com reset cadence (Antigravity windows omit durationMs). Register it in DEFAULT_RANKING_STRATEGIES so new google-antigravity sessions consult usage reports before assignment. - Regression coverage: rate-limit-utils.test.ts pins the matcher against 'Individual quota reached' and bare quota reached / quota_reached; auth-storage-antigravity-selection.test.ts proves an exhausted Gemini counter on one OAuth credential causes getApiKey to return the healthy sibling and that a less-pressured account is preferred when neither is exhausted. Fixes #2198 --- packages/ai/CHANGELOG.md | 9 + packages/ai/src/rate-limit-utils.ts | 2 +- packages/ai/src/usage/google-antigravity.ts | 44 ++-- ...auth-storage-antigravity-selection.test.ts | 204 ++++++++++++++++++ packages/ai/test/rate-limit-utils.test.ts | 19 ++ 5 files changed, 249 insertions(+), 29 deletions(-) create mode 100644 packages/ai/test/auth-storage-antigravity-selection.test.ts diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 2b10ce11b..046d8f9e7 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -120,6 +120,15 @@ - Fixed adaptive-only Claude models (Opus 4.6+, Sonnet 4.6+, Fable/Mythos 5) returning HTTP 400 `"thinking.type.disabled" is not supported for this model` whenever thinking was turned off (utility calls and forced-tool turns route through the disable path). These models accept only `thinking.type: "adaptive"`; the request builder now omits the thinking field and pins the lowest adaptive effort instead of emitting `type: "disabled"`. - Widened the OpenAI-completions first-event watchdog floor from 120s to 300s for DeepSeek V4 reasoning models hosted on the official DeepSeek API. The reasoner emits no SSE bytes until its private chain-of-thought finishes, which routinely takes longer than the generic 100s first-event budget under load — every chat then aborted with `OpenAI completions stream timed out while waiting for the first event` and silently retried. Mirrors the existing GLM coding-plan widening ([#2177](https://github.com/can1357/oh-my-pi/issues/2177)). +### Added + +- Added `antigravityRankingStrategy` and registered it for `google-antigravity` in `DEFAULT_RANKING_STRATEGIES`, so new sessions are routed to OAuth credentials with quota headroom (lowest-`remainingFraction` counter as primary, second-lowest as secondary, 24h `windowDefaults` matching `daily-cloudcode-pa.googleapis.com` resets). Without it, the existing `antigravityUsageProvider` data never reached credential selection. ([#2198](https://github.com/can1357/oh-my-pi/issues/2198)) + +### Fixed + +- Fixed `isUsageLimitError` missing Antigravity / Cloud Code Assist's `Individual quota reached` 429 phrasing. The `USAGE_LIMIT_PATTERN` only knew `quota.?exceeded` / `limit_reached`, so `auth-retry` and `AuthStorage.markUsageLimitReached` treated the response as a terminal provider error and pinned sessions to the exhausted OAuth account instead of rotating to a sibling credential. The pattern now also matches `quota.?reached`. ([#2198](https://github.com/can1357/oh-my-pi/issues/2198)) + + ## [15.10.8] - 2026-06-09 ### Added diff --git a/packages/ai/src/rate-limit-utils.ts b/packages/ai/src/rate-limit-utils.ts index 84db0cb1b..f3babc5c1 100644 --- a/packages/ai/src/rate-limit-utils.ts +++ b/packages/ai/src/rate-limit-utils.ts @@ -94,7 +94,7 @@ export function calculateRateLimitBackoffMs(reason: RateLimitReason): number { /** Detect usage/quota limit errors in error messages (persistent, requires credential switch). */ const USAGE_LIMIT_PATTERN = - /usage.?limit|usage_limit_reached|usage_not_included|limit_reached|quota.?exceeded|resource.?exhausted|exhausted your capacity|quota will reset/i; + /usage.?limit|usage_limit_reached|usage_not_included|limit_reached|quota.?exceeded|quota.?reached|resource.?exhausted|exhausted your capacity|quota will reset/i; export function isUsageLimitError(errorMessage: string): boolean { return USAGE_LIMIT_PATTERN.test(errorMessage) || ACCOUNT_RATE_LIMIT_PATTERN.test(errorMessage); diff --git a/packages/ai/src/usage/google-antigravity.ts b/packages/ai/src/usage/google-antigravity.ts index 5e935a075..8078d6ac8 100644 --- a/packages/ai/src/usage/google-antigravity.ts +++ b/packages/ai/src/usage/google-antigravity.ts @@ -301,38 +301,26 @@ export const antigravityUsageProvider: UsageProvider = { supports: params => params.provider === "google-antigravity", }; -const ANTIGRAVITY_DAILY_WINDOW_MS = 24 * 60 * 60 * 1000; +const ONE_DAY_MS = 24 * 60 * 60 * 1000; /** - * Credential ranking strategy for `google-antigravity`. Drives proactive - * multi-account selection in {@link AuthStorage} by reading the per-counter - * Antigravity usage reports. - * - * Antigravity reports one {@link UsageLimit} per backend counter (Google / - * Anthropic / OpenAI) per tier per window, and {@link fetchAntigravityUsage} - * sorts them ascending by `remainingFraction` — so `limits[0]` is always the - * most-pressured counter for the credential, and `limits[1]` (when present) - * is the next-most-pressured counter. - * - * `AuthStorage` compares the `secondary*` ranking metrics before `primary*` - * because other providers model a long-window budget as secondary. Antigravity - * does not expose a short/long split; every counter is a sibling bottleneck. - * Therefore the most-pressured counter goes in `secondary`, with the runner-up - * in `primary`, so proactive account selection always ranks the bottleneck - * before any healthier sibling counter. - * - * The Antigravity API exposes `resetTime` but not window duration, so the - * drain-rate calculation depends on `windowDefaults`. Antigravity quotas are - * effectively daily; 24h is the right fallback for both axes — any 5h tier - * still ranks correctly because both credentials are normalised against the - * same fallback. + * Antigravity quotas reset daily and are returned per backend counter + * (Anthropic / Google / OpenAI) without a fixed "primary vs secondary" + * split. `fetchAntigravityUsage` already sorts `limits` ascending by + * `remainingFraction`, so the most-pressured counter is index 0 and the + * next-most-pressured (if any) is index 1. Treat those as the windows + * AuthStorage compares across credentials — that surfaces an exhausted + * Gemini counter on one credential even when a sibling Claude counter is + * healthy, which is what was masking quota-exhausted accounts before. */ export const antigravityRankingStrategy: CredentialRankingStrategy = { findWindowLimits(report) { - return { primary: report.limits[1], secondary: report.limits[0] }; - }, - windowDefaults: { - primaryMs: ANTIGRAVITY_DAILY_WINDOW_MS, - secondaryMs: ANTIGRAVITY_DAILY_WINDOW_MS, + const primary = report.limits[0]; + const secondary = report.limits.find((limit, index) => index > 0 && limit !== primary); + return { primary, secondary }; }, + // Antigravity windows omit `durationMs`; the endpoint is + // `daily-cloudcode-pa.googleapis.com`, so fall back to 24h when computing + // drain rate. + windowDefaults: { primaryMs: ONE_DAY_MS, secondaryMs: ONE_DAY_MS }, }; diff --git a/packages/ai/test/auth-storage-antigravity-selection.test.ts b/packages/ai/test/auth-storage-antigravity-selection.test.ts new file mode 100644 index 000000000..87101e6cd --- /dev/null +++ b/packages/ai/test/auth-storage-antigravity-selection.test.ts @@ -0,0 +1,204 @@ +/** + * Antigravity OAuth ranking smoke test. Proves the + * `antigravityRankingStrategy` is wired into `DEFAULT_RANKING_STRATEGIES` + * (issue #2198): a credential whose usage report shows an exhausted + * counter must be skipped in favour of a healthy sibling on the next + * `getApiKey` call. + * + * Without the registration `getApiKey` would round-robin between + * credentials and could pin a session to the exhausted account. + */ +import { afterEach, beforeEach, describe, expect, test, vi } from "bun:test"; +import * as fs from "node:fs/promises"; +import * as os from "node:os"; +import * as path from "node:path"; +import { type AuthCredentialStore, AuthStorage, SqliteAuthCredentialStore } from "@oh-my-pi/pi-ai/auth-storage"; +import * as oauthUtils from "@oh-my-pi/pi-ai/registry/oauth"; +import type { OAuthCredentials } from "@oh-my-pi/pi-ai/registry/oauth/types"; +import type { UsageLimit, UsageProvider, UsageReport } from "@oh-my-pi/pi-ai/usage"; + +const HOUR_MS = 60 * 60 * 1000; + +type AntigravityWindowSpec = { + counter: "google" | "anthropic" | "openai" | "default"; + usedFraction: number; + resetInMs: number; +}; + +function createAntigravityLimit(spec: AntigravityWindowSpec, projectId: string): UsageLimit { + const used = Math.min(Math.max(spec.usedFraction, 0), 1); + return { + id: `google-antigravity:${spec.counter}:default:WINDOW_DAILY`, + label: `Usage (${spec.counter})`, + scope: { + provider: "google-antigravity", + projectId, + windowId: "WINDOW_DAILY", + }, + window: { + id: "WINDOW_DAILY", + label: "Default", + resetsAt: Date.now() + spec.resetInMs, + }, + amount: { + unit: "percent", + used: used * 100, + limit: 100, + remaining: (1 - used) * 100, + usedFraction: used, + remainingFraction: 1 - used, + }, + status: used >= 1 ? "exhausted" : used >= 0.9 ? "warning" : "ok", + }; +} + +function createAntigravityReport(args: { + projectId: string; + accountId: string; + windows: AntigravityWindowSpec[]; +}): UsageReport { + // fetchAntigravityUsage sorts ascending by remainingFraction; mirror + // that here so the strategy sees the same shape it would in production. + const limits = args.windows + .map(w => createAntigravityLimit(w, args.projectId)) + .sort((a, b) => (a.amount.remainingFraction ?? 1) - (b.amount.remainingFraction ?? 1)); + return { + provider: "google-antigravity", + fetchedAt: Date.now(), + limits, + metadata: { accountId: args.accountId, projectId: args.projectId }, + }; +} + +function createCredential(accountId: string, projectId: string, email: string): OAuthCredentials { + return { + access: `access-${accountId}`, + refresh: `refresh-${accountId}`, + expires: Date.now() + HOUR_MS, + accountId, + projectId, + email, + }; +} + +describe("AuthStorage google-antigravity oauth ranking", () => { + let tempDir = ""; + let store: AuthCredentialStore | null = null; + let authStorage: AuthStorage | null = null; + const usageByAccount = new Map(); + + const usageProvider: UsageProvider = { + id: "google-antigravity", + async fetchUsage(params) { + const accountId = params.credential.accountId; + if (!accountId) return null; + return usageByAccount.get(accountId) ?? null; + }, + }; + + beforeEach(async () => { + tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "pi-ai-auth-antigravity-selection-")); + store = await SqliteAuthCredentialStore.open(path.join(tempDir, "agent.db")); + authStorage = new AuthStorage(store, { + usageProviderResolver: provider => (provider === "google-antigravity" ? usageProvider : undefined), + }); + usageByAccount.clear(); + vi.spyOn(oauthUtils, "getOAuthApiKey").mockImplementation(async (_provider, credentials) => { + const credential = credentials["google-antigravity"] as OAuthCredentials | undefined; + if (!credential?.accountId) return null; + return { + apiKey: `api-${credential.accountId}`, + newCredentials: credential, + }; + }); + }); + + afterEach(async () => { + vi.restoreAllMocks(); + store?.close(); + store = null; + authStorage = null; + if (tempDir) { + await fs.rm(tempDir, { recursive: true, force: true }); + tempDir = ""; + } + }); + + test("skips antigravity account whose Gemini counter is exhausted", async () => { + if (!authStorage) throw new Error("test setup failed"); + + await authStorage.set("google-antigravity", [ + { type: "oauth", ...createCredential("acct-exhausted", "proj-exhausted", "exhausted@example.com") }, + { type: "oauth", ...createCredential("acct-healthy", "proj-healthy", "healthy@example.com") }, + ]); + + // Exhausted account: Gemini counter at 100%, Claude counter healthy. + // Pre-fix the credential was rotatable only on response-side errors, + // so a session could still be assigned to it on first use. + usageByAccount.set( + "acct-exhausted", + createAntigravityReport({ + accountId: "acct-exhausted", + projectId: "proj-exhausted", + windows: [ + { counter: "google", usedFraction: 1, resetInMs: 12 * HOUR_MS }, + { counter: "anthropic", usedFraction: 0.2, resetInMs: 12 * HOUR_MS }, + ], + }), + ); + usageByAccount.set( + "acct-healthy", + createAntigravityReport({ + accountId: "acct-healthy", + projectId: "proj-healthy", + windows: [ + { counter: "google", usedFraction: 0.3, resetInMs: 20 * HOUR_MS }, + { counter: "anthropic", usedFraction: 0.1, resetInMs: 20 * HOUR_MS }, + ], + }), + ); + + const apiKey = await authStorage.getApiKey("google-antigravity", "session-antigravity-exhausted"); + expect(apiKey).toBe("api-acct-healthy"); + }); + + test("prefers less-pressured antigravity account when neither is exhausted", async () => { + if (!authStorage) throw new Error("test setup failed"); + + await authStorage.set("google-antigravity", [ + { type: "oauth", ...createCredential("acct-loaded", "proj-loaded", "loaded@example.com") }, + { type: "oauth", ...createCredential("acct-fresh", "proj-fresh", "fresh@example.com") }, + ]); + + usageByAccount.set( + "acct-loaded", + createAntigravityReport({ + accountId: "acct-loaded", + projectId: "proj-loaded", + windows: [{ counter: "google", usedFraction: 0.8, resetInMs: 4 * HOUR_MS }], + }), + ); + usageByAccount.set( + "acct-fresh", + createAntigravityReport({ + accountId: "acct-fresh", + projectId: "proj-fresh", + windows: [{ counter: "google", usedFraction: 0.05, resetInMs: 4 * HOUR_MS }], + }), + ); + + // Sample several sessions; the weighted picker must favour the fresh + // account by a clear margin even though both are unblocked. + const counts = new Map(); + for (let i = 0; i < 60; i += 1) { + const apiKey = await authStorage.getApiKey("google-antigravity", `session-antigravity-fresh-${i}`); + if (!apiKey) continue; + counts.set(apiKey, (counts.get(apiKey) ?? 0) + 1); + } + + const fresh = counts.get("api-acct-fresh") ?? 0; + const loaded = counts.get("api-acct-loaded") ?? 0; + expect(fresh).toBeGreaterThan(loaded); + }); + +}); diff --git a/packages/ai/test/rate-limit-utils.test.ts b/packages/ai/test/rate-limit-utils.test.ts index 44ea85770..87f833674 100644 --- a/packages/ai/test/rate-limit-utils.test.ts +++ b/packages/ai/test/rate-limit-utils.test.ts @@ -87,6 +87,25 @@ describe("isUsageLimitError", () => { ), ).toBe(true); }); + + // Antigravity / Cloud Code Assist returns this phrasing for an exhausted + // project quota; `parseRateLimitReason` already maps it to QUOTA_EXHAUSTED + // via the generic `quota` substring, but `isUsageLimitError` decides + // whether the auth layer rotates to a sibling OAuth credential, so it + // must match too — otherwise the session stays pinned to the exhausted + // account (see issue #2198). + it("detects Antigravity 'Individual quota reached' as a credential-rotatable usage limit", () => { + expect( + isUsageLimitError( + "Cloud Code Assist API error (429): Individual quota reached. Contact your administrator to enable overages.", + ), + ).toBe(true); + }); + + it("detects bare 'quota reached' phrasing", () => { + expect(isUsageLimitError("quota reached")).toBe(true); + expect(isUsageLimitError("quota_reached")).toBe(true); + }); }); describe("calculateRateLimitBackoffMs", () => {