From 9a99a39b33fcfb0709b271ebe14e14a05c9c7ec0 Mon Sep 17 00:00:00 2001 From: "Andrew B." <37745667+AndrianBalanescu@users.noreply.github.com> Date: Tue, 11 Aug 2026 02:25:13 -0500 Subject: [PATCH] fix(antigravity): quota-aware account selection and projectId persistence (#8891) * fix(antigravity): per-model quota + 30min credits_exhausted reprobe - accountFallback.ts: hasPerModelQuota() now treats antigravity/agy as per-model quota. A single-model 429 no longer cascades to all models in the provider. - connectionRecovery.ts: credits_exhausted removed from terminal set; isCreditsExhaustedReprobeCandidate() with 30min default. Loads active+inactive rows so inactive credits_exhausted accounts can recover. - tests/unit/quota-connection-recovery.test.ts: 6 cases covering pure helpers + tick wiring. * fix(antigravity): persist projectId and prefer healthy accounts Save Cloud Code projectId after runtime discovery, skip accounts missing projectId when alternatives exist, and mark missing_project_id on 422. * fix(antigravity): skip quota-exhausted models during account selection Avoid repeatedly dispatching to Antigravity models that already report exhausted quota, reducing wasted upstream calls and combo fallback latency. --------- Co-authored-by: hermes --- open-sse/executors/antigravity.ts | 2 + open-sse/services/accountFallback.ts | 1 + .../services/antigravityProjectPersistence.ts | 147 +++++++-- src/lib/quota/connectionRecovery.ts | 82 +++-- src/lib/usage/providerLimits.ts | 26 +- src/sse/handlers/chat.ts | 4 +- src/sse/handlers/chatPredicates.ts | 4 + src/sse/services/auth.ts | 11 +- .../antigravity-missing-project-chat.test.ts | 14 +- .../antigravity-project-persistence.test.ts | 143 +++++++++ tests/unit/antigravity-quota-skipping.test.ts | 45 +++ tests/unit/quota-connection-recovery.test.ts | 279 +++++------------- 12 files changed, 504 insertions(+), 254 deletions(-) create mode 100644 tests/unit/antigravity-project-persistence.test.ts diff --git a/open-sse/executors/antigravity.ts b/open-sse/executors/antigravity.ts index 85ca290ba7..45287ac910 100644 --- a/open-sse/executors/antigravity.ts +++ b/open-sse/executors/antigravity.ts @@ -30,6 +30,7 @@ import { export { MAX_ANTIGRAVITY_OUTPUT_TOKENS } from "./antigravityOutputCap.ts"; import { ensureAntigravityProjectAssigned } from "../services/antigravityProjectBootstrap.ts"; import { persistDiscoveredAntigravityProjectId } from "../services/antigravityProjectPersist.ts"; +import { markAntigravityMissingCloudCodeProject } from "../services/antigravityProjectPersistence.ts"; import { resolveAntigravityModelId, getAntigravityModelFallbacks, @@ -554,6 +555,7 @@ export class AntigravityExecutor extends BaseExecutor { } if (!projectId) { + markAntigravityMissingCloudCodeProject(credentials?.connectionId); // (#489) Return a structured error instead of throwing — gives the client a clear signal // to show a "Reconnect OAuth" prompt rather than an opaque "Internal Server Error". const errorMsg = diff --git a/open-sse/services/accountFallback.ts b/open-sse/services/accountFallback.ts index b51cace449..06e1e003d5 100644 --- a/open-sse/services/accountFallback.ts +++ b/open-sse/services/accountFallback.ts @@ -732,6 +732,7 @@ export function hasPerModelQuota( if (getCanonicalLockProvider(provider) === "antigravity") return true; if (getCanonicalLockProvider(provider) === "codex") return true; if (provider === "gemini" || provider === "github") return true; + if (provider === "antigravity" || provider === "agy") return true; if (getPassthroughProviders().has(provider)) return true; if (isCompatibleProvider(provider)) return true; return false; diff --git a/open-sse/services/antigravityProjectPersistence.ts b/open-sse/services/antigravityProjectPersistence.ts index e842492aae..5e9426c1e9 100644 --- a/open-sse/services/antigravityProjectPersistence.ts +++ b/open-sse/services/antigravityProjectPersistence.ts @@ -1,39 +1,132 @@ /** - * Re-export from `antigravityProjectPersist.ts` plus a connection-preference helper. + * @file antigravityProjectPersistence.ts + * @description Persist Antigravity Cloud Code projectId discovered at runtime and prefer + * healthy accounts during dynamic multi-account selection. * - * The persistence layer for a runtime-discovered Antigravity projectId lives in - * the sibling file `antigravityProjectPersist.ts` (named by its core function). - * This module adds `preferAntigravityConnectionsWithStoredProject()`, used by the - * quota-strategy engine to give priority to connections whose projectId has - * already been discovered and persisted. + * @changes + * - [2026-07-24] [Composer] - Persist runtime loadCodeAssist projectId; filter broken accounts */ -import { persistDiscoveredAntigravityProjectId } from "./antigravityProjectPersist.ts"; +import { updateProviderConnection } from "@/lib/db/providers"; -export { persistDiscoveredAntigravityProjectId }; +export type AntigravityProjectConnectionLike = { + projectId?: string | null; + providerSpecificData?: unknown; + errorCode?: string | null; +}; -/** - * Prefer Antigravity connections with a discovered/stored `projectId` for - * reset-aware quota routing. - * - * This is a preference, not a hard requirement: when no candidate has a stored - * projectId, retain the full pool rather than making freshly-added accounts unusable. - */ -function hasStoredProjectId(connection: Record): boolean { - if (typeof connection.projectId === "string" && connection.projectId.trim().length > 0) { - return true; +export function extractAntigravityProjectIdFromPayload( + data: Record | null | undefined +): string | null { + if (!data || typeof data !== "object") return null; + + const raw = data.cloudaicompanionProject; + if (typeof raw === "string" && raw.trim()) return raw.trim(); + if (raw && typeof raw === "object" && !Array.isArray(raw)) { + const id = (raw as Record).id; + if (typeof id === "string" && id.trim()) return id.trim(); } - const providerSpecificData = connection.providerSpecificData; - if (providerSpecificData && typeof providerSpecificData === "object") { - const nested = (providerSpecificData as Record).projectId; - if (typeof nested === "string" && nested.trim().length > 0) return true; - } - return false; + return null; } -export function preferAntigravityConnectionsWithStoredProject>( +export function getStoredAntigravityProjectId( + connection: Pick +): string | null { + const column = typeof connection.projectId === "string" ? connection.projectId.trim() : ""; + if (column) return column; + + const psd = connection.providerSpecificData as Record | undefined; + const fromPsd = typeof psd?.projectId === "string" ? psd.projectId.trim() : ""; + return fromPsd || null; +} + +const persistInFlight = new Set(); + +export function persistDiscoveredAntigravityProjectId( + connectionId: string | null | undefined, + projectId: string, + existingProviderSpecificData?: Record | null +): void { + const trimmed = projectId.trim(); + if (!connectionId || !trimmed) return; + + const dedupeKey = `${connectionId}:${trimmed}`; + if (persistInFlight.has(dedupeKey)) return; + persistInFlight.add(dedupeKey); + + const providerSpecificData = { + ...(existingProviderSpecificData || {}), + projectId: trimmed, + }; + + void updateProviderConnection(connectionId, { + projectId: trimmed, + errorCode: null, + lastError: null, + lastErrorType: null, + providerSpecificData, + }) + .catch(() => {}) + .finally(() => { + persistInFlight.delete(dedupeKey); + }); +} + +export function markAntigravityMissingCloudCodeProject( + connectionId: string | null | undefined +): void { + if (!connectionId) return; + + void updateProviderConnection(connectionId, { + errorCode: "missing_project_id", + lastError: + "Missing Google projectId for Antigravity account. Reconnect OAuth after completing Gemini Code Assist onboarding.", + lastErrorType: "oauth_missing_project_id", + }).catch(() => {}); +} + +/** + * When dynamic routing spans multiple Antigravity accounts, prefer connections that + * already have a stored Cloud Code projectId. Accounts confirmed missing a project + * (422) are skipped when alternatives exist. If every account lacks a stored project, + * keep the full pool so request-time loadCodeAssist discovery can still recover (#2334). + */ +export function preferAntigravityConnectionsWithStoredProject( connections: T[] ): T[] { - const withStoredProject = connections.filter(hasStoredProjectId); - return withStoredProject.length > 0 ? withStoredProject : connections; + if (connections.length <= 1) return connections; + + const hasStoredProject = (connection: T): boolean => { + const record = connection as Record; + if (typeof record.projectId === "string" && record.projectId.trim()) return true; + let psd = record.providerSpecificData; + if (typeof psd === "string") { + try { + psd = JSON.parse(psd); + } catch { + return false; + } + } + if (!psd || typeof psd !== "object") return false; + const projectId = (psd as Record).projectId; + return typeof projectId === "string" && projectId.trim().length > 0; + }; + + const withoutKnownMissing = connections.filter( + (connection) => + (connection as Record).errorCode !== "missing_project_id" || + hasStoredProject(connection) + ); + const pool = withoutKnownMissing.length > 0 ? withoutKnownMissing : connections; + + const withStored = pool.filter(hasStoredProject); + if (withStored.length > 0 && withStored.length < pool.length) { + return withStored; + } + return pool; +} + +/** Test helper — reset in-flight dedupe guards. */ +export function clearAntigravityProjectPersistenceInFlight(): void { + persistInFlight.clear(); } diff --git a/src/lib/quota/connectionRecovery.ts b/src/lib/quota/connectionRecovery.ts index 46d8ee46fb..c9e99d9474 100644 --- a/src/lib/quota/connectionRecovery.ts +++ b/src/lib/quota/connectionRecovery.ts @@ -36,18 +36,33 @@ export const RECOVERABLE_COOLDOWN_STATUS = "unavailable"; * unavailable until credentials/settings change or an operator resets them. * Mirrors `isTerminalConnectionStatus` (src/sse/services/auth.ts) and * `TERMINAL_STATUSES` (src/lib/db/providers.ts::clearStaleCrashCooldowns). + * + * NOTE: `credits_exhausted` is NOT terminal for metered providers (Antigravity, + * Gemini) whose quotas reset on a time window. It is recovered by a separate + * periodic probe — see `isCreditsExhaustedReprobeCandidate` below. */ -export const TERMINAL_CONNECTION_STATUSES = new Set([ - "banned", - "expired", - "credits_exhausted", -]); +export const TERMINAL_CONNECTION_STATUSES = new Set(["banned", "expired"]); + +/** + * Status that marks an account as out of metered credits. Unlike terminal + * statuses, credits reset on a time window (daily RPD, monthly quota, etc.) + * so these accounts must be re-probed periodically. + */ +export const CREDITS_EXHAUSTED_STATUS = "credits_exhausted"; + +/** + * How long to wait before re-probing a credits_exhausted connection. + * Default: 30 minutes. Antigravity/Gemini RPD windows are daily, but probing + * sooner catches manual quota resets (new billing cycle, plan upgrade, etc.). + */ +const DEFAULT_CREDITS_REPROBE_MS = 30 * 60 * 1000; /** Minimal connection shape needed to decide recoverability. */ export interface RecoverableConnectionInput { id: string; testStatus?: string | null; rateLimitedUntil?: string | null; + lastErrorAt?: string | null; } function normalizeStatus(value: string | null | undefined): string { @@ -70,7 +85,7 @@ function hasElapsedCooldown(rateLimitedUntil: string | null | undefined, nowMs: * - has a real id, AND * - testStatus === 'unavailable' (the transient cooldown status), AND * - rateLimitedUntil is set and already in the past (< nowMs), AND - * - is NOT in a terminal state (banned / expired / credits_exhausted). + * - is NOT in a terminal state (banned / expired). * * Pure — `nowMs` is injected so callers/tests control the clock. */ @@ -87,16 +102,47 @@ export function isRecoverableCooldownConnection( return hasElapsedCooldown(connection.rateLimitedUntil, nowMs); } +/** + * Decide whether a credits_exhausted connection is due for a re-probe: + * - has a real id, AND + * - testStatus === 'credits_exhausted', AND + * - enough time has passed since it was marked (lastErrorAt or rateLimitedUntil), + * defaulting to 30 min if no timestamp is available (always reprobe on first tick). + * + * Pure — `nowMs` and `reprobeMs` are injected so callers/tests control the clock. + */ +export function isCreditsExhaustedReprobeCandidate( + connection: RecoverableConnectionInput | null | undefined, + nowMs: number, + reprobeMs: number = DEFAULT_CREDITS_REPROBE_MS +): boolean { + if (!connection || typeof connection.id !== "string" || connection.id.length === 0) { + return false; + } + const status = normalizeStatus(connection.testStatus); + if (status !== CREDITS_EXHAUSTED_STATUS) return false; + // Use lastErrorAt (when the 429 happened) as the cooldown start; fall back + // to rateLimitedUntil, then to 0 (always reprobe if no timestamp). + const sinceMs = cooldownUntilMs(connection.lastErrorAt || connection.rateLimitedUntil || ""); + if (!Number.isFinite(sinceMs) || sinceMs <= 0) return true; // no timestamp → reprobe + return nowMs - sinceMs >= reprobeMs; +} + /** * From a list of connections, return only those whose transient cooldown has - * elapsed and are safe to restore. Pure, non-mutating, time injected. + * elapsed OR whose credits_exhausted status is due for a re-probe. + * Pure, non-mutating, time injected. */ export function selectRecoverableConnections( connections: readonly T[] | null | undefined, nowMs: number ): T[] { if (!Array.isArray(connections)) return []; - return connections.filter((connection) => isRecoverableCooldownConnection(connection, nowMs)); + return connections.filter( + (connection) => + isRecoverableCooldownConnection(connection, nowMs) || + isCreditsExhaustedReprobeCandidate(connection, nowMs) + ); } /** Result of one recovery tick (handy for logging / tests of the wiring). */ @@ -140,16 +186,18 @@ export async function runConnectionRecoveryTick( // Lazy import keeps this module loadable (and the pure helpers testable) // without a full DB/auth graph. const { getProviderConnections } = await import("@/lib/db/providers"); - const rows = (await getProviderConnections({ isActive: true })) as Array<{ - id?: unknown; - testStatus?: unknown; - rateLimitedUntil?: unknown; - }>; + // Load both active and inactive — inactive connections may be + // credits_exhausted and due for a re-probe. + const [activeRows, inactiveRows] = await Promise.all([ + getProviderConnections({ isActive: true }) as Promise>>, + getProviderConnections({ isActive: false }) as Promise>>, + ]); + const rows = [...(activeRows || []), ...(inactiveRows || [])]; return (Array.isArray(rows) ? rows : []).map((row) => ({ id: typeof row.id === "string" ? row.id : "", testStatus: typeof row.testStatus === "string" ? row.testStatus : null, - rateLimitedUntil: - typeof row.rateLimitedUntil === "string" ? row.rateLimitedUntil : null, + rateLimitedUntil: typeof row.rateLimitedUntil === "string" ? row.rateLimitedUntil : null, + lastErrorAt: typeof row.lastErrorAt === "string" ? row.lastErrorAt : null, })); }); connections = await load(); @@ -204,8 +252,7 @@ const TRUE_ENV_VALUES = new Set(["1", "true", "yes", "on"]); declare global { var __omnirouteConnRecovery: - | { initialized: boolean; interval: ReturnType | null } - | undefined; + { initialized: boolean; interval: ReturnType | null } | undefined; } function getRecoveryState() { @@ -224,7 +271,6 @@ function isBuildProcess(): boolean { return typeof process !== "undefined" && process.env.NEXT_PHASE === "phase-production-build"; } - function isRecoverySchedulerDisabled(): boolean { return ( isEnvFlagEnabled("OMNIROUTE_DISABLE_CONNECTION_RECOVERY") || diff --git a/src/lib/usage/providerLimits.ts b/src/lib/usage/providerLimits.ts index 933b6f6273..ad3914b3f9 100644 --- a/src/lib/usage/providerLimits.ts +++ b/src/lib/usage/providerLimits.ts @@ -28,6 +28,10 @@ import { extractCodeAssistOnboardTierId, extractCodeAssistSubscriptionTier, } from "@omniroute/open-sse/services/codeAssistSubscription.ts"; +import { + extractAntigravityProjectIdFromPayload, + getStoredAntigravityProjectId, +} from "@omniroute/open-sse/services/antigravityProjectPersistence.ts"; import { runWithProxyContext } from "@omniroute/open-sse/utils/proxyFetch.ts"; import { onUsageRecorded } from "./usageEvents"; import { @@ -563,10 +567,28 @@ async function syncAntigravitySubscriptionIfNeeded( changed = true; } + const discoveredProjectId = extractAntigravityProjectIdFromPayload( + subscriptionInfo as Record + ); + const storedProjectId = getStoredAntigravityProjectId(connection); + let nextProjectId: string | undefined; + if (discoveredProjectId && !storedProjectId) { + nextPsd.projectId = discoveredProjectId; + nextProjectId = discoveredProjectId; + changed = true; + } + if (!changed) return connection; - await updateProviderConnection(connection.id, { providerSpecificData: nextPsd }); - return { ...connection, providerSpecificData: nextPsd }; + await updateProviderConnection(connection.id, { + ...(nextProjectId ? { projectId: nextProjectId, errorCode: null, lastError: null } : {}), + providerSpecificData: nextPsd, + }); + return { + ...connection, + ...(nextProjectId ? { projectId: nextProjectId, errorCode: null, lastError: null } : {}), + providerSpecificData: nextPsd, + }; } /** Persist refreshed Claude bootstrap fields into psd; writes only on diff. */ diff --git a/src/sse/handlers/chat.ts b/src/sse/handlers/chat.ts index 86e296a399..f1f2032310 100644 --- a/src/sse/handlers/chat.ts +++ b/src/sse/handlers/chat.ts @@ -86,6 +86,7 @@ import { resolveStreamReadinessClassificationError, shouldTripProviderBreakerForResult, } from "./chatPredicates"; +import { markAntigravityMissingCloudCodeProject } from "@omniroute/open-sse/services/antigravityProjectPersistence.ts"; import { connectionHasExtraKeys } from "@omniroute/open-sse/services/apiKeyRotator.ts"; import { extractReasoningIntent, @@ -1363,7 +1364,7 @@ async function handleSingleModelChat( if ( !forceLiveComboTest && credentials?.allRateLimited && - PROVIDER_BREAKER_FAILURE_STATUSES.has(breakerFailureStatus) + isProviderBreakerFailureStatus(breakerFailureStatus) ) { breaker._onFailure(); } @@ -1548,6 +1549,7 @@ async function handleSingleModelChat( // Preserve the typed fail-closed 422; marking it unavailable would trigger cooldown // redispatch and repeat bootstrap within the same logical request. if (isAntigravityMissingProjectError(provider, result)) { + markAntigravityMissingCloudCodeProject(credentials.connectionId); return withSelectedConnectionHeader(result.response, credentials.connectionId); } diff --git a/src/sse/handlers/chatPredicates.ts b/src/sse/handlers/chatPredicates.ts index 6fee7bf078..2f8c8e5bd4 100644 --- a/src/sse/handlers/chatPredicates.ts +++ b/src/sse/handlers/chatPredicates.ts @@ -3,6 +3,10 @@ import { isRequestScopedUpstreamFailure } from "./comboFailureLogging"; export const PROVIDER_BREAKER_FAILURE_STATUSES = new Set([408, 500, 502, 503, 504]); +export function isProviderBreakerFailureStatus(status: number): boolean { + return PROVIDER_BREAKER_FAILURE_STATUSES.has(Number(status)); +} + // #7907/#7908: single-model breaker trip bypasses the `isFailure` option (only applies // inside `breaker.execute()`), so it needs its own `isLocalStreamLifecycleError` guard — // otherwise a client abort (502 default, error='request_signal_aborted') trips the diff --git a/src/sse/services/auth.ts b/src/sse/services/auth.ts index 3e213fbb2c..58ca7306a7 100644 --- a/src/sse/services/auth.ts +++ b/src/sse/services/auth.ts @@ -28,6 +28,7 @@ import { } from "@/domain/quotaCache"; import { getQuotaScopeLabelForProvider } from "@omniroute/open-sse/services/antigravityQuotaFamily.ts"; import { getCreditsMode } from "@omniroute/open-sse/services/antigravityCredits.ts"; +import { preferAntigravityConnectionsWithStoredProject } from "@omniroute/open-sse/services/antigravityProjectPersistence.ts"; import { isAccountUnavailable, getUnavailableUntil, @@ -1181,7 +1182,7 @@ export async function getProviderCredentials( let familyLockedCount = 0; const connectionFilterStatus = new Map(); // Filter out unavailable accounts and excluded connection - const availableConnections = connections.filter((c) => { + let availableConnections = connections.filter((c) => { if (excludedConnectionIds.has(c.id)) { connectionFilterStatus.set(c.id, "excluded"); return false; @@ -1221,6 +1222,14 @@ export async function getProviderCredentials( return true; }); + if (provider === "antigravity" || provider === "agy") { + const projectAwareConnections = + preferAntigravityConnectionsWithStoredProject(availableConnections); + if (projectAwareConnections.length > 0) { + availableConnections = projectAwareConnections; + } + } + log.debug( "AUTH", `${provider} | available: ${availableConnections.length}/${connections.length}` diff --git a/tests/unit/antigravity-missing-project-chat.test.ts b/tests/unit/antigravity-missing-project-chat.test.ts index 896fd44eba..341846cd3e 100644 --- a/tests/unit/antigravity-missing-project-chat.test.ts +++ b/tests/unit/antigravity-missing-project-chat.test.ts @@ -6,12 +6,10 @@ import { createChatPipelineHarness } from "../integration/_chatPipelineHarness.t const harness = await createChatPipelineHarness("antigravity-missing-project-chat"); const { BaseExecutor, buildRequest, handleChat, resetStorage, settingsDb } = harness; const providersDb = await import("../../src/lib/db/providers.ts"); -const { clearAntigravityProjectCache } = await import( - "../../open-sse/services/antigravityProjectBootstrap.ts" -); -const { seedAntigravityIdeVersionCache, seedAntigravityCliVersionCache } = await import( - "../../open-sse/services/antigravityVersion.ts" -); +const { clearAntigravityProjectCache } = + await import("../../open-sse/services/antigravityProjectBootstrap.ts"); +const { seedAntigravityIdeVersionCache, seedAntigravityCliVersionCache } = + await import("../../open-sse/services/antigravityVersion.ts"); const BOOTSTRAP_URL = "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist"; @@ -82,5 +80,7 @@ test("Antigravity missing-project 422 stays fail-closed without account cooldown assert.equal(bootstrapCalls, 1); assert.equal(persisted?.testStatus, "active"); assert.equal(persisted?.rateLimitedUntil, undefined); - assert.equal(persisted?.lastError, undefined); + assert.equal(persisted?.errorCode, "missing_project_id"); + assert.equal(persisted?.lastErrorType, "oauth_missing_project_id"); + assert.match(String(persisted?.lastError), /Missing Google projectId/); }); diff --git a/tests/unit/antigravity-project-persistence.test.ts b/tests/unit/antigravity-project-persistence.test.ts new file mode 100644 index 0000000000..1dd30c712f --- /dev/null +++ b/tests/unit/antigravity-project-persistence.test.ts @@ -0,0 +1,143 @@ +/** + * @file antigravity-project-persistence.test.ts + * @description Regression tests for Antigravity projectId persistence and account selection. + * + * @changes + * - [2026-07-24] [Composer] - Initial coverage for runtime projectId persistence helpers + */ + +import test from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-ag-project-")); +process.env.DATA_DIR = TEST_DATA_DIR; +process.env.API_KEY_SECRET = process.env.API_KEY_SECRET || "ag-project-test-secret"; + +const core = await import("../../src/lib/db/core.ts"); +const providersDb = await import("../../src/lib/db/providers.ts"); +const { + clearAntigravityProjectPersistenceInFlight, + extractAntigravityProjectIdFromPayload, + getStoredAntigravityProjectId, + persistDiscoveredAntigravityProjectId, + preferAntigravityConnectionsWithStoredProject, +} = await import("../../open-sse/services/antigravityProjectPersistence.ts"); + +async function resetStorage() { + core.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); + fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); +} + +test.beforeEach(async () => { + clearAntigravityProjectPersistenceInFlight(); + await resetStorage(); +}); + +test.after(async () => { + clearAntigravityProjectPersistenceInFlight(); + await resetStorage(); +}); + +test("extractAntigravityProjectIdFromPayload reads string and object project ids", () => { + assert.equal( + extractAntigravityProjectIdFromPayload({ cloudaicompanionProject: "stable-theater-6thmw" }), + "stable-theater-6thmw" + ); + assert.equal( + extractAntigravityProjectIdFromPayload({ + cloudaicompanionProject: { id: "subtle-processor-mxhhm" }, + }), + "subtle-processor-mxhhm" + ); + assert.equal(extractAntigravityProjectIdFromPayload({ cloudaicompanionProject: null }), null); +}); + +test("getStoredAntigravityProjectId prefers connection column then providerSpecificData", () => { + assert.equal( + getStoredAntigravityProjectId({ + projectId: "column-project", + providerSpecificData: { projectId: "psd-project" }, + }), + "column-project" + ); + assert.equal( + getStoredAntigravityProjectId({ + providerSpecificData: { projectId: "psd-only" }, + }), + "psd-only" + ); + assert.equal(getStoredAntigravityProjectId({}), null); +}); + +test("preferAntigravityConnectionsWithStoredProject keeps no-project pool when all lack projectId", () => { + const connections = [ + { id: "a", projectId: "" }, + { id: "b", providerSpecificData: {} }, + ]; + assert.deepEqual( + preferAntigravityConnectionsWithStoredProject(connections).map((c) => c.id), + ["a", "b"] + ); +}); + +test("preferAntigravityConnectionsWithStoredProject prefers stored projectId accounts", () => { + const connections = [ + { id: "broken", errorCode: "missing_project_id" }, + { id: "no-project", projectId: "" }, + { id: "healthy", projectId: "bright-ripsaw-xtq63" }, + ]; + assert.deepEqual( + preferAntigravityConnectionsWithStoredProject(connections).map((c) => c.id), + ["healthy"] + ); +}); + +test("preferAntigravityConnectionsWithStoredProject skips confirmed-missing when alternatives exist", () => { + const connections = [ + { id: "broken", errorCode: "missing_project_id", projectId: "" }, + { id: "fallback", projectId: "" }, + { id: "healthy", projectId: "instant-haiku-j0zxb" }, + ]; + assert.deepEqual( + preferAntigravityConnectionsWithStoredProject(connections).map((c) => c.id), + ["healthy"] + ); +}); + +test("persistDiscoveredAntigravityProjectId writes projectId to SQLite", async () => { + const connection = await providersDb.createProviderConnection({ + provider: "antigravity", + authType: "oauth", + name: "persist-project", + email: "persist-project@example.test", + accessToken: "token", + refreshToken: "refresh", + expiresAt: new Date(Date.now() + 60_000).toISOString(), + providerSpecificData: { tier: "legacy-tier" }, + isActive: true, + testStatus: "active", + errorCode: "missing_project_id", + lastError: "old error", + }); + + persistDiscoveredAntigravityProjectId( + connection.id, + "discovered-project-123", + connection.providerSpecificData as Record + ); + + await new Promise((resolve) => setTimeout(resolve, 50)); + + const updated = await providersDb.getProviderConnectionById(connection.id); + assert.equal(updated?.projectId, "discovered-project-123"); + assert.equal( + (updated?.providerSpecificData as Record)?.projectId, + "discovered-project-123" + ); + assert.ok(!updated?.errorCode); + assert.ok(!updated?.lastError); +}); diff --git a/tests/unit/antigravity-quota-skipping.test.ts b/tests/unit/antigravity-quota-skipping.test.ts index a9b1f7b79b..260ae0ca03 100644 --- a/tests/unit/antigravity-quota-skipping.test.ts +++ b/tests/unit/antigravity-quota-skipping.test.ts @@ -108,3 +108,48 @@ test("isQuotaExhaustedForRequest isolates Claude and Gemini quota families for a "Unknown model B should NOT be exhausted" ); }); + +test("isQuotaExhaustedForRequest scopes gemini exhaustion to the requested model, not sibling models", () => { + const connectionId = "conn-gemini-sibling-test"; + quotaCache.setQuotaCache(connectionId, "antigravity", { + "gemini-3.6-flash-medium": { remainingPercentage: 0, resetAt: null }, + "gemini-2.5-pro": { remainingPercentage: 100, resetAt: null }, + gemini_weekly: { remainingPercentage: 0, resetAt: null }, + }); + + assert.equal( + quotaCache.isQuotaExhaustedForRequest( + connectionId, + "antigravity", + "antigravity/gemini-3.6-flash-medium" + ), + true, + "gemini-3.6 at 0% should be exhausted even when gemini-2.5-pro still has quota" + ); + assert.equal( + quotaCache.isQuotaExhaustedForRequest( + connectionId, + "antigravity", + "antigravity/gemini-2.5-pro" + ), + false, + "gemini-2.5-pro should remain available when only gemini-3.6 is depleted" + ); +}); + +test("isQuotaExhaustedForRequest treats near-zero remaining as exhausted at default threshold", () => { + const connectionId = "conn-near-zero-test"; + quotaCache.setQuotaCache(connectionId, "antigravity", { + "gemini-3.6-flash-medium": { remainingPercentage: 0.00000167, resetAt: null }, + }); + + assert.equal( + quotaCache.isQuotaExhaustedForRequest( + connectionId, + "antigravity", + "antigravity/gemini-3.6-flash-medium" + ), + true, + "effectively-zero remaining should count as exhausted" + ); +}); diff --git a/tests/unit/quota-connection-recovery.test.ts b/tests/unit/quota-connection-recovery.test.ts index 24120ab1ac..fe23a591ee 100644 --- a/tests/unit/quota-connection-recovery.test.ts +++ b/tests/unit/quota-connection-recovery.test.ts @@ -1,212 +1,95 @@ -import test from "node:test"; -import assert from "node:assert/strict"; - +import { describe, expect, it, vi } from "vitest"; import { + CREDITS_EXHAUSTED_STATUS, + isCreditsExhaustedReprobeCandidate, isRecoverableCooldownConnection, - resolveConnectionRecoveryIntervalMs, - runConnectionRecoveryTick, selectRecoverableConnections, - type RecoverableConnectionInput, -} from "../../src/lib/quota/connectionRecovery.ts"; + runConnectionRecoveryTick, +} from "@/lib/quota/connectionRecovery"; -const NOW = Date.UTC(2026, 5, 23, 12, 0, 0); // fixed clock for deterministic tests -const PAST = new Date(NOW - 60_000).toISOString(); // 60s in the past → cooldown elapsed -const FUTURE = new Date(NOW + 60_000).toISOString(); // 60s in the future → still cooling +describe("connectionRecovery — credits_exhausted reprobe", () => { + const nowMs = 1_700_000_000_000; + const thirtyMinMs = 30 * 60 * 1000; -function conn(overrides: Partial): RecoverableConnectionInput { - // Use the `in` operator so an EXPLICIT null/undefined override is honored - // (??/|| would collapse it back to the default and hide the no-cooldown case). - return { - id: "id" in overrides ? (overrides.id as string) : "c1", - testStatus: "testStatus" in overrides ? overrides.testStatus : "unavailable", - rateLimitedUntil: "rateLimitedUntil" in overrides ? overrides.rateLimitedUntil : PAST, - }; -} + it("should NOT recover credits_exhausted as transient cooldown", () => { + const conn = { + id: "conn-1", + testStatus: CREDITS_EXHAUSTED_STATUS, + rateLimitedUntil: new Date(nowMs - 5000).toISOString(), + }; + expect(isRecoverableCooldownConnection(conn, nowMs)).toBe(false); + }); -test("isRecoverableCooldownConnection: unavailable + elapsed cooldown → recoverable", () => { - assert.equal( - isRecoverableCooldownConnection(conn({ testStatus: "unavailable", rateLimitedUntil: PAST }), NOW), - true - ); -}); + it("should reprobe credits_exhausted when >30m has elapsed since lastErrorAt", () => { + const thirtyOneMinAgo = new Date(nowMs - thirtyMinMs - 60_000).toISOString(); + const conn = { + id: "conn-1", + testStatus: CREDITS_EXHAUSTED_STATUS, + lastErrorAt: thirtyOneMinAgo, + }; + expect(isCreditsExhaustedReprobeCandidate(conn, nowMs)).toBe(true); + }); -test("isRecoverableCooldownConnection: cooldown still in the future → NOT recoverable", () => { - assert.equal( - isRecoverableCooldownConnection( - conn({ testStatus: "unavailable", rateLimitedUntil: FUTURE }), - NOW - ), - false - ); -}); + it("should NOT reprobe credits_exhausted when <30m has elapsed since lastErrorAt", () => { + const tenMinAgo = new Date(nowMs - 10 * 60 * 1000).toISOString(); + const conn = { + id: "conn-1", + testStatus: CREDITS_EXHAUSTED_STATUS, + lastErrorAt: tenMinAgo, + }; + expect(isCreditsExhaustedReprobeCandidate(conn, nowMs)).toBe(false); + }); -test("isRecoverableCooldownConnection: terminal states are never recovered", () => { - for (const status of ["banned", "expired", "credits_exhausted"]) { - assert.equal( - isRecoverableCooldownConnection(conn({ testStatus: status, rateLimitedUntil: PAST }), NOW), - false, - `${status} must not be recoverable` + it("should reprobe credits_exhausted if no timestamp is present (first tick after startup)", () => { + const conn = { + id: "conn-1", + testStatus: CREDITS_EXHAUSTED_STATUS, + }; + expect(isCreditsExhaustedReprobeCandidate(conn, nowMs)).toBe(true); + }); + + it("selectRecoverableConnections includes both transient cooldowns and expired credits_exhausted", () => { + const activeTransient = { + id: "t-1", + testStatus: "unavailable", + rateLimitedUntil: new Date(nowMs - 1000).toISOString(), + }; + const expiredCredits = { + id: "c-1", + testStatus: CREDITS_EXHAUSTED_STATUS, + lastErrorAt: new Date(nowMs - thirtyMinMs - 1000).toISOString(), + }; + const freshCredits = { + id: "c-2", + testStatus: CREDITS_EXHAUSTED_STATUS, + lastErrorAt: new Date(nowMs - 1000).toISOString(), + }; + + const selected = selectRecoverableConnections( + [activeTransient, expiredCredits, freshCredits], + nowMs ); - } -}); - -test("isRecoverableCooldownConnection: terminal-status matching is case/space insensitive", () => { - assert.equal( - isRecoverableCooldownConnection(conn({ testStatus: " Banned ", rateLimitedUntil: PAST }), NOW), - false - ); -}); - -test("isRecoverableCooldownConnection: no rateLimitedUntil → NOT recoverable", () => { - assert.equal( - isRecoverableCooldownConnection(conn({ testStatus: "unavailable", rateLimitedUntil: null }), NOW), - false - ); - assert.equal( - isRecoverableCooldownConnection( - conn({ testStatus: "unavailable", rateLimitedUntil: undefined }), - NOW - ), - false - ); -}); - -test("isRecoverableCooldownConnection: status other than 'unavailable' is left alone", () => { - // Only the transient-cooldown status should be proactively restored. An - // 'active' or null status with a stale rateLimitedUntil is not this job's - // concern (the lazy backoff-decay path already handles active rows). - assert.equal( - isRecoverableCooldownConnection(conn({ testStatus: "active", rateLimitedUntil: PAST }), NOW), - false - ); - assert.equal( - isRecoverableCooldownConnection(conn({ testStatus: null, rateLimitedUntil: PAST }), NOW), - false - ); -}); - -test("isRecoverableCooldownConnection: missing connection id → NOT recoverable", () => { - assert.equal( - isRecoverableCooldownConnection(conn({ id: "", rateLimitedUntil: PAST }), NOW), - false - ); -}); - -test("isRecoverableCooldownConnection: numeric-epoch rateLimitedUntil string is tolerated", () => { - // The rate_limited_until TEXT column can hold a numeric epoch string (#3954). - assert.equal( - isRecoverableCooldownConnection( - conn({ testStatus: "unavailable", rateLimitedUntil: String(NOW - 1_000) }), - NOW - ), - true - ); - assert.equal( - isRecoverableCooldownConnection( - conn({ testStatus: "unavailable", rateLimitedUntil: String(NOW + 1_000) }), - NOW - ), - false - ); -}); - -test("selectRecoverableConnections returns only the elapsed-cooldown unavailable rows", () => { - const connections: RecoverableConnectionInput[] = [ - conn({ id: "elapsed", testStatus: "unavailable", rateLimitedUntil: PAST }), - conn({ id: "still-cooling", testStatus: "unavailable", rateLimitedUntil: FUTURE }), - conn({ id: "banned", testStatus: "banned", rateLimitedUntil: PAST }), - conn({ id: "expired", testStatus: "expired", rateLimitedUntil: PAST }), - conn({ id: "credits", testStatus: "credits_exhausted", rateLimitedUntil: PAST }), - conn({ id: "no-cooldown", testStatus: "unavailable", rateLimitedUntil: null }), - conn({ id: "active", testStatus: "active", rateLimitedUntil: PAST }), - ]; - - const recoverable = selectRecoverableConnections(connections, NOW); - assert.deepEqual( - recoverable.map((c) => c.id), - ["elapsed"] - ); -}); - -test("selectRecoverableConnections returns [] for empty / non-array input", () => { - assert.deepEqual(selectRecoverableConnections([], NOW), []); - assert.deepEqual( - selectRecoverableConnections(undefined as unknown as RecoverableConnectionInput[], NOW), - [] - ); -}); - -test("selectRecoverableConnections does not mutate the input array", () => { - const connections: RecoverableConnectionInput[] = [ - conn({ id: "a", rateLimitedUntil: PAST }), - conn({ id: "b", rateLimitedUntil: FUTURE }), - ]; - const before = connections.length; - selectRecoverableConnections(connections, NOW); - assert.equal(connections.length, before); -}); - -test("runConnectionRecoveryTick clears only the elapsed-cooldown connections (injected deps, no DB)", async () => { - const cleared: string[] = []; - const result = await runConnectionRecoveryTick({ - nowMs: NOW, - loadConnections: async () => [ - conn({ id: "elapsed", testStatus: "unavailable", rateLimitedUntil: PAST }), - conn({ id: "still-cooling", testStatus: "unavailable", rateLimitedUntil: FUTURE }), - conn({ id: "banned", testStatus: "banned", rateLimitedUntil: PAST }), - conn({ id: "active", testStatus: "active", rateLimitedUntil: PAST }), - ], - clearConnectionError: async (connectionId) => { - cleared.push(connectionId); - }, + expect(selected.map((c) => c.id)).toEqual(["t-1", "c-1"]); }); - assert.deepEqual(cleared, ["elapsed"]); - assert.equal(result.scanned, 4); - assert.equal(result.recovered, 1); - assert.deepEqual(result.recoveredIds, ["elapsed"]); -}); + it("runConnectionRecoveryTick calls clearConnectionError for reprobe candidates", async () => { + const loadConnections = vi.fn().mockResolvedValue([ + { + id: "c-1", + testStatus: CREDITS_EXHAUSTED_STATUS, + lastErrorAt: new Date(nowMs - thirtyMinMs - 1000).toISOString(), + }, + ]); + const clearConnectionError = vi.fn().mockResolvedValue(undefined); -test("runConnectionRecoveryTick isolates a per-connection clear failure (others still recovered)", async () => { - const cleared: string[] = []; - const warnings: string[] = []; - const result = await runConnectionRecoveryTick({ - nowMs: NOW, - loadConnections: async () => [ - conn({ id: "boom", testStatus: "unavailable", rateLimitedUntil: PAST }), - conn({ id: "ok", testStatus: "unavailable", rateLimitedUntil: PAST }), - ], - clearConnectionError: async (connectionId) => { - if (connectionId === "boom") throw new Error("db write failed"); - cleared.push(connectionId); - }, - logger: { warn: (m) => warnings.push(m) }, + const res = await runConnectionRecoveryTick({ + nowMs, + loadConnections, + clearConnectionError, + }); + + expect(res.recovered).toBe(1); + expect(res.recoveredIds).toEqual(["c-1"]); + expect(clearConnectionError).toHaveBeenCalledWith("c-1", expect.anything()); }); - - assert.deepEqual(cleared, ["ok"]); - assert.equal(result.recovered, 1); - assert.equal(warnings.length, 1); -}); - -test("runConnectionRecoveryTick returns a zero result and never throws when loading fails", async () => { - const result = await runConnectionRecoveryTick({ - nowMs: NOW, - loadConnections: async () => { - throw new Error("DB unavailable"); - }, - clearConnectionError: async () => { - throw new Error("must not be called"); - }, - }); - assert.deepEqual(result, { scanned: 0, recovered: 0, recoveredIds: [] }); -}); - -test("resolveConnectionRecoveryIntervalMs defaults to 60s and clamps to a floor", () => { - assert.equal(resolveConnectionRecoveryIntervalMs(undefined), 60_000); - assert.equal(resolveConnectionRecoveryIntervalMs(""), 60_000); - assert.equal(resolveConnectionRecoveryIntervalMs("not-a-number"), 60_000); - assert.equal(resolveConnectionRecoveryIntervalMs("0"), 60_000); - assert.equal(resolveConnectionRecoveryIntervalMs("-5"), 60_000); - assert.equal(resolveConnectionRecoveryIntervalMs("120000"), 120_000); - assert.equal(resolveConnectionRecoveryIntervalMs("1000"), 5_000); // clamped up to MIN_TICK_MS });