Skip to content
Open
22 changes: 11 additions & 11 deletions package-lock.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

46 changes: 43 additions & 3 deletions src/plugin.ts
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ import { AccountManager, type ModelFamily, parseRateLimitReason, calculateBackof
import { createAutoUpdateCheckerHook } from "./hooks/auto-update-checker";
import { loadConfig, initRuntimeConfig, type AntigravityConfig } from "./plugin/config";
import { createSessionRecoveryHook, getRecoverySuccessToast } from "./plugin/recovery";
import { checkAccountsQuota } from "./plugin/quota";
import { checkAccountsQuota, triggerAsyncQuotaRefreshForAll } from "./plugin/quota";
import { initDiskSignatureCache } from "./plugin/cache";
import { createProactiveRefreshQueue, type ProactiveRefreshQueue } from "./plugin/refresh-queue";
import { initLogger, createLogger } from "./plugin/logger";
Expand Down Expand Up @@ -90,7 +90,7 @@ let softQuotaToastShown = false;
let rateLimitToastShown = false;

// Module-level reference to AccountManager for access from auth.login
let activeAccountManager: import("./plugin/accounts").AccountManager | null = null;
let activeAccountManager: AccountManager | null = null;

function cleanupToastCooldowns(): void {
if (rateLimitToastCooldowns.size > MAX_TOAST_COOLDOWN_ENTRIES) {
Expand Down Expand Up @@ -2022,7 +2022,23 @@ export const createAntigravityPlugin = (providerId: string) => async (
tokenConsumed = getTokenTracker().consume(account.index);
}

const response = await fetch(prepared.request, prepared.init);
// Check if we should proactively refresh all quotas
if (accountManager.shouldRefreshAllQuotas()) {
pushDebug("proactive-quota-refresh: pool is mostly blocked, refreshing all");
void triggerAsyncQuotaRefreshForAll(accountManager, client, providerId);
}

// Create a combined signal for timeout and user abort
const timeoutMs = (config.request_timeout_seconds ?? 180) * 1000;
const timeoutSignal = AbortSignal.timeout(timeoutMs);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

CRITICAL: AbortSignal.timeout() compatibility issue - will cause runtime errors

AbortSignal.timeout() was added in Node.js 20.11.0, but this project supports Node.js >=20.0.0. Users on Node.js 20.0.0 - 20.10.x will encounter a TypeError: AbortSignal.timeout is not a function runtime error.

Suggested change
const timeoutSignal = AbortSignal.timeout(timeoutMs);
// Create a combined signal for timeout and user abort
const timeoutMs = (config.request_timeout_seconds ?? 180) * 1000;
const timeoutSignal = (typeof AbortSignal.timeout === 'function')
? AbortSignal.timeout(timeoutMs)
: createTimeoutSignal(timeoutMs);
const combinedSignal = abortSignal
? (AbortSignal as any).any([abortSignal, timeoutSignal])
: timeoutSignal;

You'll need to add a createTimeoutSignal polyfill function alongside the existing AbortSignal.any() polyfill.

const combinedSignal = abortSignal
? (AbortSignal as any).any([abortSignal, timeoutSignal])
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

CRITICAL: AbortSignal.any() compatibility issue - may cause runtime errors

AbortSignal.any() was added in Node.js 20.17.0, but this project supports Node.js >=20.0.0. Users on Node.js 20.0.0 - 20.16.x will encounter a runtime error.

Suggested change
? (AbortSignal as any).any([abortSignal, timeoutSignal])
// Create a combined signal for timeout and user abort
const timeoutMs = (config.request_timeout_seconds ?? 180) * 1000;
const timeoutSignal = AbortSignal.timeout(timeoutMs);
const combinedSignal = abortSignal
? (abortSignal.any ? abortSignal.any([abortSignal, timeoutSignal]) : mergeAbortSignals(abortSignal, timeoutSignal))
: timeoutSignal;

Alternatively, you could add a polyfill at the top of the file:

// Polyfill for older Node.js versions
if (typeof AbortSignal.any === 'undefined') {
  AbortSignal.any = function(signals: AbortSignal[]): AbortSignal {
    const controller = new AbortController();
    for (const signal of signals) {
      if (signal.aborted) {
        controller.abort(signal.reason);
        break;
      }
      signal.addEventListener('abort', () => {
        controller.abort(signal.reason);
      }, { once: true });
    }
    return controller.signal;
  };
}

: timeoutSignal;

const response = await fetch(prepared.request, {
...prepared.init,
signal: combinedSignal
});
pushDebug(`status=${response.status} ${response.statusText}`);


Expand Down Expand Up @@ -2403,6 +2419,30 @@ export const createAntigravityPlugin = (providerId: string) => async (
tokenConsumed = false;
}

// [CRITICAL] Check for AbortError from user vs timeout
if (error instanceof Error && (error.name === "AbortError" || error.name === "TimeoutError")) {
if (abortSignal?.aborted) {
// User pressed ESC - stop everything immediately to prevent spin loop and memory leak
pushDebug("user-interrupted: stopping request loop");
throw error;
}

// This was a request timeout (stuck account)
const timeoutSec = config.request_timeout_seconds ?? 180;
pushDebug(`request-timeout: account ${account.index} stuck for ${timeoutSec}s, rotating`);
getHealthTracker().recordFailure(account.index);
accountManager.markAccountCoolingDown(account, 60000, "network-error");

await showToast(
`⏳ Account stuck (${timeoutSec}s). Rotating to next available...`,
"warning"
);

shouldSwitchAccount = true;
lastError = error;
break;
}

// Handle recoverable thinking errors - retry with forced recovery
if (error instanceof Error && error.message === "THINKING_RECOVERY_NEEDED") {
// Only retry once with forced recovery to avoid infinite loops
Expand Down
20 changes: 20 additions & 0 deletions src/plugin/accounts.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1160,9 +1160,29 @@ export class AccountManager {
if (account) {
account.cachedQuota = quotaGroups;
account.cachedQuotaUpdatedAt = nowMs();
this.requestSaveToDisk();
}
}

/**
* Heuristic to determine if we should refresh all quotas in the pool.
* Triggers when the pool is mostly blocked or unhealthy.
*/
shouldRefreshAllQuotas(): boolean {
const enabled = this.getEnabledAccounts();
if (enabled.length === 0) return false;

const now = nowMs();
// Refresh if more than 75% of enabled accounts are marked as rate-limited or over soft quota
const blockedCount = enabled.filter(acc => {
const isRateLimited = Object.values(acc.rateLimitResetTimes).some(t => t !== undefined && t > now);
const isCoolingDown = acc.coolingDownUntil !== undefined && acc.coolingDownUntil > now;
return isRateLimited || isCoolingDown;
}).length;

return (blockedCount / enabled.length) >= 0.75;
}

isAccountOverSoftQuota(account: ManagedAccount, family: ModelFamily, thresholdPercent: number, cacheTtlMs: number, model?: string | null): boolean {
return isOverSoftQuotaThreshold(account, family, thresholdPercent, cacheTtlMs, model);
}
Expand Down
10 changes: 10 additions & 0 deletions src/plugin/config/schema.ts
Original file line number Diff line number Diff line change
Expand Up @@ -267,6 +267,15 @@ export const AntigravityConfigSchema = z.object({
* @default 300 (5 minutes)
*/
max_rate_limit_wait_seconds: z.number().min(0).max(3600).default(300),

/**
* Maximum time in seconds for a single API request before timing out.
* When a request times out, the account is marked unhealthy and the
* plugin automatically rotates to the next available account.
*
* @default 180 (3 minutes)
*/
request_timeout_seconds: z.number().min(30).max(1800).default(180),

/**
* @deprecated Kept only for backward compatibility.
Expand Down Expand Up @@ -464,6 +473,7 @@ export const DEFAULT_CONFIG: AntigravityConfig = {
proactive_refresh_buffer_seconds: 1800,
proactive_refresh_check_interval_seconds: 300,
max_rate_limit_wait_seconds: 300,
request_timeout_seconds: 180,
quota_fallback: false,
cli_first: false,
account_selection_strategy: 'hybrid',
Expand Down
35 changes: 35 additions & 0 deletions src/plugin/quota.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,11 @@ import { refreshAccessToken } from "./token";
import { getModelFamily } from "./transform/model-resolver";
import type { PluginClient, OAuthAuthDetails } from "./types";
import type { AccountMetadataV3 } from "./storage";
import type { AccountManager } from "./accounts";

const FETCH_TIMEOUT_MS = 10000;
const POOL_REFRESH_COOLDOWN_MS = 60000; // 1 minute
let lastPoolRefreshTime = 0;

export type QuotaGroup = "claude" | "gemini-pro" | "gemini-flash";

Expand Down Expand Up @@ -393,3 +396,35 @@ export async function checkAccountsQuota(
logQuotaFetch("complete", accounts.length, `ok=${results.filter(r => r.status === "ok").length} errors=${results.filter(r => r.status === "error").length}`);
return results;
}

/**
* Proactively refreshes quotas for all accounts in the background.
* Updates the account manager with new quota data.
*/
export async function triggerAsyncQuotaRefreshForAll(
accountManager: AccountManager,
client: PluginClient,
providerId: string,
): Promise<void> {
const now = Date.now();
if (now - lastPoolRefreshTime < POOL_REFRESH_COOLDOWN_MS) {
return;
}
lastPoolRefreshTime = now;

try {
const accountsMetadata = accountManager.getAccountsForQuotaCheck();
if (accountsMetadata.length === 0) return;

const results = await checkAccountsQuota(accountsMetadata, client, providerId);

for (const result of results) {
if (result.status === "ok" && result.quota) {
accountManager.updateQuotaCache(result.index, result.quota.groups);
}
}
} catch (error) {
// Proactive refresh is best-effort - log and ignore
console.error("[ProactiveQuota] Failed to refresh all quotas", error);
}
}