Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/adapters/base.ts
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,8 @@ export interface AdapterRequest {
export interface AdapterFetchContext {
/** Remains attached to the returned response body after the response headers arrive. */
abortSignal?: AbortSignal;
/** OAuth account identity used for provider-local cooldown bookkeeping. */
accountId?: string;
/** Deadline for receiving response headers on each attempt, not for consuming the response body. */
timeoutMs?: number;
/** Return final non-2xx responses untouched so the caller can own the error-body read. */
Expand Down
24 changes: 24 additions & 0 deletions src/adapters/google-antigravity-hosts.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
const DAILY_ANTIGRAVITY_HOST = "https://daily-cloudcode-pa.googleapis.com";
const PROD_ANTIGRAVITY_HOST = "https://cloudcode-pa.googleapis.com";

/**
* Return the configured Antigravity endpoint followed by its daily/production peer.
* The configured value is preserved so tests and future pinned environments keep their
* explicit first choice; the fallback is always one of Google's two known hosts.
*/
export function antigravityHostCandidates(configuredBase: string): string[] {
const configured = configuredBase.replace(/\/+$/, "");
const other = configured === DAILY_ANTIGRAVITY_HOST
? PROD_ANTIGRAVITY_HOST
: DAILY_ANTIGRAVITY_HOST;
return [...new Set([configured, other])];
}

/** OAuth bearer requests must not use a cleartext host, even if generic baseUrl config allows http. */
export function isAntigravityHttpsHost(host: string): boolean {
try {
return new URL(host).protocol === "https:";
} catch {
return false;
}
}
7 changes: 7 additions & 0 deletions src/adapters/google-errors.ts
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,12 @@ function googleErrorDetail(payloadText: string): { message?: string; status?: st
};
}

const ANTIGRAVITY_GEO_BLOCKED_MARKER = "user location is not supported for the api use";

export function isAntigravityGeoBlockedBody(payloadText: string): boolean {
return payloadText.toLowerCase().includes(ANTIGRAVITY_GEO_BLOCKED_MARKER);
}

function classifyGoogle(label: string, status: number | undefined, enumStatus: string | undefined, text: string): string {
const lower = `${enumStatus ?? ""} ${text}`.toLowerCase();
const quotaExhausted =
Expand All @@ -29,6 +35,7 @@ function classifyGoogle(label: string, status: number | undefined, enumStatus: s
if (status === 401 || enumStatus === "UNAUTHENTICATED" || lower.includes("unauthenticated") || lower.includes("invalid authentication") || lower.includes("expired")) {
return `${label} authentication failed`;
}
if (isAntigravityGeoBlockedBody(lower)) return `${label} location not supported`;
if (status === 403 || enumStatus === "PERMISSION_DENIED" || lower.includes("permission denied") || lower.includes("access denied")) {
return `${label} access denied`;
}
Expand Down
39 changes: 38 additions & 1 deletion src/adapters/google-http.ts
Original file line number Diff line number Diff line change
@@ -1,7 +1,13 @@
import type { AdapterFetchContext, AdapterRequest } from "./base";
import { isQuotaExhaustedBody, retryableGoogleStatus, safeGoogleHttpErrorMessage } from "./google-errors";
import {
isAntigravityGeoBlockedBody,
isQuotaExhaustedBody,
retryableGoogleStatus,
safeGoogleHttpErrorMessage,
} from "./google-errors";
import { repairGoogleInvalidRequestBody } from "./google-wire-compiler";
import { normalizeUpstreamHttpErrorResponse, readDisplaySafeErrorPayloadText } from "./upstream-http-error";
import { recordAntigravityCooldown } from "../oauth/antigravity-routing";
import {
abortError,
cancelResponseBodyBestEffort,
Expand All @@ -26,6 +32,34 @@ async function normalizeFinalGoogleError(label: string, res: Response, signal?:
});
}

function retryAfterMs(value: string | null, now = Date.now()): number | undefined {
const text = value?.trim();
if (!text) return undefined;
if (/^\d+(?:\.\d+)?$/.test(text)) {
const seconds = Number(text);
return Number.isFinite(seconds) && seconds > 0 ? Math.ceil(seconds * 1000) : undefined;
}
const timestamp = Date.parse(text);
return Number.isFinite(timestamp) && timestamp > now ? timestamp - now : undefined;
}

async function recordAntigravityHttpCooldown(
response: Response,
accountId: string | undefined,
): Promise<void> {
if (!accountId || (response.status !== 429 && response.status !== 403)) return;
const payloadText = await readDisplaySafeErrorPayloadText(response.clone());
if (response.status === 429) {
recordAntigravityCooldown(
accountId,
isQuotaExhaustedBody(payloadText) ? "quota_exhausted" : "rate_limited",
retryAfterMs(response.headers.get("retry-after")),
);
} else if (isAntigravityGeoBlockedBody(payloadText)) {
recordAntigravityCooldown(accountId, "geo_blocked");
}
}

/**
* Fetch a Google-family upstream with Kiro-style hardening: per-attempt timeout
* (`AbortSignal.any([parent, timeout])`), bounded retry on transient status / network errors,
Expand Down Expand Up @@ -53,6 +87,9 @@ export async function fetchGoogleWithRetry(
headers: activeRequest.headers,
body: activeRequest.body,
}, timeoutMs, ctx.abortSignal, ctx.stream, executor);
if (label === "Antigravity") {
await recordAntigravityHttpCooldown(res, ctx.accountId);
}
if (res.status === 400 && repairInvalid400 && !compatibilityReplayUsed) {
let payloadText = "";
try {
Expand Down
2 changes: 2 additions & 0 deletions src/lib/state-store-registrations.ts
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ import {
sweepExpiredXaiPermanentFailureVerdicts,
} from "../oauth";
import { sweepExpiredAnthropicRoutingHealth } from "../oauth/anthropic-routing";
import { sweepExpiredAntigravityRoutingHealth } from "../oauth/antigravity-routing";
import { listLiveOAuthAccountKeys, reconcileOAuthReauthState } from "../oauth/store";
import { reconcileGuardianBackoff } from "../oauth/token-guardian";
import { sweepExpiredApiKeyCooldowns } from "../providers/key-failover";
Expand Down Expand Up @@ -83,6 +84,7 @@ export const STATE_STORE_REGISTRATIONS = [
reconcileGeneration: reconcileComboTargetCooldowns,
},
{ name: "anthropic-routing-health", sweepExpired: sweepExpiredAnthropicRoutingHealth },
{ name: "antigravity-routing-health", sweepExpired: sweepExpiredAntigravityRoutingHealth },
{ name: "xai-refresh-verdicts", sweepExpired: sweepExpiredXaiPermanentFailureVerdicts },
{ name: "responses-continuation", sweepExpired: sweepExpiredResponseStates },
{ name: "antigravity-replay", sweepExpired: sweepExpiredAntigravityReplay },
Expand Down
132 changes: 132 additions & 0 deletions src/oauth/antigravity-routing.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,132 @@
export type AntigravityCooldownReason = "rate_limited" | "quota_exhausted" | "geo_blocked";

const DEFAULT_RATE_LIMITED_COOLDOWN_MS = 5_000;
const MAX_RATE_LIMITED_COOLDOWN_MS = 60_000;
const DEFAULT_QUOTA_EXHAUSTED_COOLDOWN_MS = 24 * 60 * 60_000;
const MAX_QUOTA_EXHAUSTED_COOLDOWN_MS = 7 * 24 * 60 * 60_000;
const GEO_BLOCKED_COOLDOWN_MS = 24 * 60 * 60_000;

type AntigravityAccountHealth = {
cooldownUntil: number;
};

const accountHealth = new Map<string, AntigravityAccountHealth>();

function positiveDurationOrDefault(
durationMs: number | undefined,
defaultMs: number,
maxMs?: number,
): number {
if (typeof durationMs !== "number" || !Number.isFinite(durationMs) || durationMs <= 0) {
return defaultMs;
}
return maxMs === undefined ? durationMs : Math.min(durationMs, maxMs);
}

function cooldownDurationMs(
reason: AntigravityCooldownReason,
retryAfterMs: number | undefined,
): number {
switch (reason) {
case "rate_limited":
return positiveDurationOrDefault(
retryAfterMs,
DEFAULT_RATE_LIMITED_COOLDOWN_MS,
MAX_RATE_LIMITED_COOLDOWN_MS,
);
case "quota_exhausted":
return positiveDurationOrDefault(
retryAfterMs,
DEFAULT_QUOTA_EXHAUSTED_COOLDOWN_MS,
MAX_QUOTA_EXHAUSTED_COOLDOWN_MS,
);
case "geo_blocked":
return GEO_BLOCKED_COOLDOWN_MS;
}
}

export function recordAntigravityCooldown(
accountId: string,
reason: AntigravityCooldownReason,
retryAfterMs?: number,
now = Date.now(),
): void {
const cooldownUntil = now + cooldownDurationMs(reason, retryAfterMs);
const current = accountHealth.get(accountId);
if (!current || current.cooldownUntil < cooldownUntil) {
accountHealth.set(accountId, { cooldownUntil });
}
}

export function isAntigravityAccountInCooldown(accountId: string, now = Date.now()): boolean {
const health = accountHealth.get(accountId);
if (!health) return false;
if (health.cooldownUntil <= now) {
accountHealth.delete(accountId);
return false;
}
return true;
}

export function nextAntigravityAccount(
accountIds: string[],
activeId: string | undefined,
now = Date.now(),
): string | undefined {
if (accountIds.length === 0) return undefined;

const activeIndex = activeId === undefined ? -1 : accountIds.indexOf(activeId);
const startIndex = activeIndex < 0 ? 0 : activeIndex + 1;
for (let offset = 0; offset < accountIds.length; offset += 1) {
const accountId = accountIds[(startIndex + offset) % accountIds.length]!;
if (activeId !== undefined && accountId === activeId) continue;
if (!isAntigravityAccountInCooldown(accountId, now)) return accountId;
}
return undefined;
}

export function sweepExpiredAntigravityRoutingHealth(now = Date.now()): number {
let removed = 0;
for (const [accountId, health] of accountHealth) {
if (health.cooldownUntil > now) continue;
accountHealth.delete(accountId);
removed += 1;
}
return removed;
}

export function clearAntigravityAccountCooldown(accountId: string): void {
accountHealth.delete(accountId);
}

export const ANTIGRAVITY_MISSING_PROJECT_MESSAGE =
"Antigravity requires a discovered Cloud Code Assist project id (re-run `ocx login google-antigravity`).";

export type BindAntigravityProjectFailure = {
ok: false;
status: 400;
type: "invalid_request_error";
message: string;
};

export type BindAntigravityProjectSuccess<T extends { project?: string }> = {
ok: true;
provider: T & { project: string };
};

/** Pair Cloud Code Assist `project` with the credential in use. Never keep a previous account's id. */
export function bindAntigravityProject<T extends { project?: string }>(
provider: T,
projectId: string | undefined,
): BindAntigravityProjectSuccess<T> | BindAntigravityProjectFailure {
const project = typeof projectId === "string" ? projectId.trim() : "";
if (!project) {
return {
ok: false,
status: 400,
type: "invalid_request_error",
message: ANTIGRAVITY_MISSING_PROJECT_MESSAGE,
};
}
return { ok: true, provider: { ...provider, project } };
}
Loading
Loading