ampcode-connector 0.1.19 → 0.1.21

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "ampcode-connector",
-  "version": "0.1.19",
+  "version": "0.1.21",
   "description": "Proxy AmpCode through local OAuth subscriptions (Claude Code, Codex, Gemini CLI, Antigravity)",
   "license": "MIT",
   "repository": {

package/src/cli/status.ts

@@ -21,13 +21,13 @@ export interface ProviderStatus {
 const PROVIDERS: { name: ProviderName; label: string; sublabel?: string }[] = [
   { name: "anthropic", label: "Claude Code" },
   { name: "codex", label: "OpenAI Codex" },
-  { name: "google", label: "Google", sublabel: "Gemini CLI + Antigravity" },
+  { name: "google", label: "Google" },
 ];
 
 const POOL_MAP: Record<ProviderName, QuotaPool[]> = {
   anthropic: ["anthropic"],
   codex: ["codex"],
-  google: ["gemini", "antigravity"],
+  google: ["google"],
 };
 
 function connectionOf(name: ProviderName, account: number, creds: Credentials): ConnectionStatus {

package/src/constants.ts

@@ -1,12 +1,12 @@
 /** Single source of truth — no magic strings scattered across files. */
 
 export const CODE_ASSIST_ENDPOINT = "https://cloudcode-pa.googleapis.com";
-export const ANTIGRAVITY_DAILY_ENDPOINT = "https://daily-cloudcode-pa.sandbox.googleapis.com";
+export const ANTIGRAVITY_DAILY_ENDPOINT = "https://daily-cloudcode-pa.googleapis.com";
+export const ANTIGRAVITY_DAILY_SANDBOX_ENDPOINT = "https://daily-cloudcode-pa.sandbox.googleapis.com";
 export const AUTOPUSH_ENDPOINT = "https://autopush-cloudcode-pa.sandbox.googleapis.com";
 export const DEFAULT_ANTIGRAVITY_PROJECT = "rising-fact-p41fc";
 
 export const ANTHROPIC_API_URL = "https://api.anthropic.com";
-export const OPENAI_API_URL = "https://api.openai.com";
 export const CODEX_BASE_URL = "https://chatgpt.com/backend-api";
 
 /** Codex-specific headers required by the ChatGPT backend. */
@@ -18,11 +18,13 @@ export const codexHeaders = {
   CONVERSATION_ID: "conversation_id",
 } as const;
 
+export const CODEX_CLI_VERSION = "0.101.0";
+
 export const codexHeaderValues = {
   BETA_RESPONSES: "responses=experimental",
   ORIGINATOR: "codex_cli_rs",
-  VERSION: "0.101.0",
-  USER_AGENT: `codex_cli_rs/0.101.0 (${process.platform} ${process.arch})`,
+  VERSION: CODEX_CLI_VERSION,
+  USER_AGENT: `codex_cli_rs/${CODEX_CLI_VERSION} (${process.platform} ${process.arch})`,
 } as const;
 
 /** Map Amp CLI paths → ChatGPT backend paths.
@@ -40,18 +42,6 @@ export const GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token";
 export const TOKEN_EXPIRY_BUFFER_MS = 5 * 60 * 1000;
 export const CLAUDE_CODE_VERSION = "2.1.77";
 
-export const stainlessHeaders: Readonly<Record<string, string>> = {
-  "X-Stainless-Helper-Method": "stream",
-  "X-Stainless-Retry-Count": "0",
-  "X-Stainless-Runtime-Version": "v24.3.0",
-  "X-Stainless-Package-Version": "0.74.0",
-  "X-Stainless-Runtime": "node",
-  "X-Stainless-Lang": "js",
-  "X-Stainless-Arch": process.arch,
-  "X-Stainless-Os": process.platform === "darwin" ? "MacOS" : process.platform === "win32" ? "Windows" : "Linux",
-  "X-Stainless-Timeout": "600",
-};
-
 export const claudeCodeBetas = [
   "claude-code-20250219",
   "oauth-2025-04-20",

package/src/providers/forward.ts

@@ -1,6 +1,7 @@
 /** HTTP forwarding with transport-level retry, SSE proxying, and response rewriting. */
 
 import { logger } from "../utils/logger.ts";
+import { apiError } from "../utils/responses.ts";
 import * as sse from "../utils/streaming.ts";
 
 export interface ForwardOptions {
@@ -17,6 +18,24 @@ const RETRYABLE_STATUS = new Set([408, 500, 502, 503, 504]);
 const MAX_RETRIES = 3;
 const RETRY_DELAY_MS = 500;
 
+const PASSTHROUGH_HEADERS = [
+  "content-type",
+  "retry-after",
+  "x-request-id",
+  "x-ratelimit-limit-requests",
+  "x-ratelimit-remaining-requests",
+  "x-ratelimit-reset-requests",
+];
+
+function copyHeaders(source: Headers): Headers {
+  const dest = new Headers();
+  for (const name of PASSTHROUGH_HEADERS) {
+    const value = source.get(name);
+    if (value !== null) dest.set(name, value);
+  }
+  return dest;
+}
+
 export async function forward(opts: ForwardOptions): Promise<Response> {
   for (let attempt = 0; attempt <= MAX_RETRIES; attempt++) {
     let response: Response;
@@ -34,7 +53,7 @@ export async function forward(opts: ForwardOptions): Promise<Response> {
         await Bun.sleep(RETRY_DELAY_MS * (attempt + 1));
         continue;
       }
-      throw err;
+      return transportErrorResponse(opts.providerName, err);
     }
 
     // Retry on server errors (429 handled at routing layer)
@@ -70,21 +89,25 @@ export async function forward(opts: ForwardOptions): Promise<Response> {
         });
       }
 
+      const headers = copyHeaders(response.headers);
+      headers.set("Content-Type", "application/json");
       return new Response(errorBody, {
         status: response.status,
-        headers: { "Content-Type": "application/json" },
+        headers,
       });
     }
 
     const isSSE = contentType.includes("text/event-stream") || opts.streaming;
     if (isSSE) return sse.proxy(response, opts.rewrite);
 
+    const headers = copyHeaders(response.headers);
+
     if (opts.rewrite) {
       const text = await response.text();
-      return new Response(opts.rewrite(text), { status: response.status, headers: { "Content-Type": contentType } });
+      return new Response(opts.rewrite(text), { status: response.status, headers });
     }
 
-    return new Response(response.body, { status: response.status, headers: { "Content-Type": contentType } });
+    return new Response(response.body, { status: response.status, headers });
   }
 
   // Unreachable, but TypeScript needs it
@@ -92,5 +115,32 @@ export async function forward(opts: ForwardOptions): Promise<Response> {
 }
 
 export function denied(providerName: string): Response {
-  return Response.json({ error: `No ${providerName} OAuth token available. Run login first.` }, { status: 401 });
+  return apiError(401, `No ${providerName} OAuth token available. Run login first.`);
+}
+
+function transportErrorResponse(providerName: string, err: unknown): Response {
+  const message = transportErrorMessage(providerName, err);
+  logger.error(`${providerName} transport error after retries exhausted`, { error: String(err) });
+  return apiError(502, message, "connection_error");
+}
+
+function transportErrorMessage(providerName: string, err: unknown): string {
+  const base = `${providerName} connection error after retries were exhausted.`;
+  const details = String(err);
+
+  if (providerName !== "Anthropic") {
+    return `${base} ${details}`;
+  }
+
+  const looksLikeReset =
+    details.includes("ECONNRESET") ||
+    details.includes("socket connection was closed unexpectedly") ||
+    details.includes("tls") ||
+    details.includes("network");
+
+  if (!looksLikeReset) {
+    return `${base} ${details}`;
+  }
+
+  return `${base} ${details} This is often a local network issue rather than an OAuth bug: check Wi-Fi MTU (1492 is a common fix), hotspot stability, and iPhone dual-SIM Cellular Data Switching if you are tethering.`;
 }

package/src/providers/google.ts

@@ -0,0 +1,324 @@
+/** Unified Google provider — merges Gemini CLI and Antigravity strategies
+ *  with internal fallback. Tries preferred strategy first, then falls back. */
+
+import { google as config } from "../auth/configs.ts";
+import * as oauth from "../auth/oauth.ts";
+import * as store from "../auth/store.ts";
+import { ANTIGRAVITY_DAILY_ENDPOINT, ANTIGRAVITY_DAILY_SANDBOX_ENDPOINT, CODE_ASSIST_ENDPOINT } from "../constants.ts";
+import { buildUrl, maybeWrap, withUnwrap } from "../utils/code-assist.ts";
+import { logger } from "../utils/logger.ts";
+import * as path from "../utils/path.ts";
+import { apiError } from "../utils/responses.ts";
+import type { Provider } from "./base.ts";
+import { denied, forward } from "./forward.ts";
+
+const GOOGLE_CLIENT_METADATA = JSON.stringify({
+  ideType: "IDE_UNSPECIFIED",
+  platform: "PLATFORM_UNSPECIFIED",
+  pluginType: "GEMINI",
+});
+
+interface GoogleStrategy {
+  name: string;
+  headers: Readonly<Record<string, string>>;
+  endpoints: readonly string[];
+  modelMapper?: (model: string) => string;
+  wrapOpts: {
+    userAgent: "antigravity" | "pi-coding-agent";
+    requestIdPrefix: "agent" | "pi";
+    requestType?: "agent" | "image_gen";
+  };
+}
+
+const geminiStrategy: GoogleStrategy = {
+  name: "gemini",
+  headers: {
+    "User-Agent": "google-cloud-sdk vscode_cloudshelleditor/0.1",
+    "X-Goog-Api-Client": "gl-node/22.17.0",
+    "Client-Metadata": GOOGLE_CLIENT_METADATA,
+  },
+  endpoints: [CODE_ASSIST_ENDPOINT],
+  wrapOpts: {
+    userAgent: "pi-coding-agent",
+    requestIdPrefix: "pi",
+  },
+};
+
+/** Antigravity uses different model names than what Amp CLI sends. */
+const antigravityModelMap: Record<string, string> = {
+  "gemini-3-flash-preview": "gemini-3-flash",
+  "gemini-3-pro-preview": "gemini-3-pro-high",
+  "gemini-3-pro-image-preview": "gemini-3.1-flash-image",
+  "gemini-3.1-flash-image-preview": "gemini-3.1-flash-image",
+};
+
+const antigravityStrategy: GoogleStrategy = {
+  name: "antigravity",
+  headers: {
+    "User-Agent": "antigravity/1.104.0 darwin/arm64",
+    "X-Goog-Api-Client": "google-cloud-sdk vscode_cloudshelleditor/0.1",
+    "Client-Metadata": GOOGLE_CLIENT_METADATA,
+  },
+  endpoints: [ANTIGRAVITY_DAILY_ENDPOINT, ANTIGRAVITY_DAILY_SANDBOX_ENDPOINT, CODE_ASSIST_ENDPOINT],
+  modelMapper: (model: string) => antigravityModelMap[model] ?? model,
+  wrapOpts: {
+    userAgent: "antigravity",
+    requestIdPrefix: "agent",
+    requestType: "agent",
+  },
+};
+
+const strategies: readonly GoogleStrategy[] = [geminiStrategy, antigravityStrategy];
+
+/** Models that only work on the antigravity strategy. */
+const ANTIGRAVITY_ONLY_MODELS = new Set(["gemini-3-pro-image-preview", "gemini-3.1-flash-image-preview"]);
+
+const COOLDOWN_MS = 60_000;
+
+interface StrategyPreference {
+  strategy: GoogleStrategy;
+  until: number;
+}
+
+// Per-account strategy preference: after success, prefer that strategy;
+// after failure, skip it for COOLDOWN_MS.
+const preferredStrategy = new Map<number, StrategyPreference>();
+const cooldowns = new Map<string, number>(); // key: `${account}:${strategy.name}`
+
+function cooldownKey(account: number, strategy: GoogleStrategy): string {
+  return `${account}:${strategy.name}`;
+}
+
+function getOrderedStrategies(account: number, model?: string): GoogleStrategy[] {
+  const now = Date.now();
+
+  // Some models only work on antigravity — skip other strategies entirely
+  if (model && ANTIGRAVITY_ONLY_MODELS.has(model)) {
+    const cd = cooldowns.get(cooldownKey(account, antigravityStrategy));
+    if (cd && cd > now) return [];
+    return [antigravityStrategy];
+  }
+
+  const pref = preferredStrategy.get(account);
+  const ordered =
+    pref && pref.until > now ? [pref.strategy, ...strategies.filter((s) => s !== pref.strategy)] : [...strategies];
+
+  return ordered.filter((s) => {
+    const cd = cooldowns.get(cooldownKey(account, s));
+    return !cd || cd <= now;
+  });
+}
+
+function markSuccess(account: number, strategy: GoogleStrategy): void {
+  preferredStrategy.set(account, { strategy, until: Date.now() + COOLDOWN_MS * 10 });
+  cooldowns.delete(cooldownKey(account, strategy));
+}
+
+function markFailure(account: number, strategy: GoogleStrategy): void {
+  cooldowns.set(cooldownKey(account, strategy), Date.now() + COOLDOWN_MS);
+  const pref = preferredStrategy.get(account);
+  if (pref?.strategy === strategy) {
+    preferredStrategy.delete(account);
+  }
+}
+
+/** Buffer an SSE response and merge all chunks into a single JSON response.
+ *  Used when we force streamGenerateContent but the client expects non-streaming JSON.
+ *  Accumulates all candidate parts across chunks (image inlineData may be in earlier chunks). */
+async function bufferSSEToJSON(response: Response): Promise<Response | null> {
+  const text = await response.text();
+  const chunks: Record<string, unknown>[] = [];
+
+  for (const line of text.split("\n")) {
+    if (!line.startsWith("data: ")) continue;
+    const data = line.slice(6).trim();
+    if (!data || data === "[DONE]") continue;
+    try {
+      chunks.push(JSON.parse(data) as Record<string, unknown>);
+    } catch {
+      // skip non-JSON
+    }
+  }
+
+  if (chunks.length === 0) return null;
+  if (chunks.length === 1) {
+    return new Response(JSON.stringify(chunks[0]), {
+      status: 200,
+      headers: { "Content-Type": "application/json" },
+    });
+  }
+
+  // Merge: accumulate parts from all chunks into the first candidate
+  const merged = chunks[0]! as Record<string, unknown>;
+  const allParts: unknown[] = [];
+
+  for (const chunk of chunks) {
+    const candidates = chunk.candidates as { content?: { parts?: unknown[] } }[] | undefined;
+    if (!candidates) continue;
+    for (const candidate of candidates) {
+      if (candidate.content?.parts) {
+        allParts.push(...candidate.content.parts);
+      }
+    }
+  }
+
+  // Use last chunk's metadata (finishReason, usageMetadata)
+  const last = chunks[chunks.length - 1]! as Record<string, unknown>;
+  const lastCandidates = last.candidates as Record<string, unknown>[] | undefined;
+  const mergedCandidates = merged.candidates as Record<string, unknown>[] | undefined;
+
+  if (mergedCandidates?.[0] && allParts.length > 0) {
+    const content = (mergedCandidates[0] as Record<string, unknown>).content as Record<string, unknown> | undefined;
+    if (content) content.parts = allParts;
+    if (lastCandidates?.[0]) {
+      (mergedCandidates[0] as Record<string, unknown>).finishReason = (
+        lastCandidates[0] as Record<string, unknown>
+      ).finishReason;
+    }
+  }
+  if (last.usageMetadata) merged.usageMetadata = last.usageMetadata;
+
+  return new Response(JSON.stringify(merged), {
+    status: 200,
+    headers: { "Content-Type": "application/json" },
+  });
+}
+
+/** Generate a fallback project ID when none is stored (matches CLIProxyAPI behavior). */
+function generateProjectId(): string {
+  const adjectives = ["useful", "bright", "swift", "calm", "bold"];
+  const nouns = ["fuze", "wave", "spark", "flow", "core"];
+  const adj = adjectives[Math.floor(Math.random() * adjectives.length)]!;
+  const noun = nouns[Math.floor(Math.random() * nouns.length)]!;
+  const rand = crypto.randomUUID().slice(0, 5).toLowerCase();
+  return `${adj}-${noun}-${rand}`;
+}
+
+export const provider: Provider = {
+  name: "Google",
+  routeDecision: "LOCAL_GOOGLE",
+
+  isAvailable: (account?: number) =>
+    account !== undefined ? !!store.get("google", account)?.refreshToken : oauth.ready(config),
+
+  accountCount: () => oauth.accountCount(config),
+
+  async forward(sub, body, _originalHeaders, rewrite, account = 0) {
+    const accessToken = await oauth.token(config, account);
+    if (!accessToken) return denied("Google");
+
+    const creds = store.get("google", account);
+    const projectId = creds?.projectId || generateProjectId();
+    const email = creds?.email;
+
+    const modelAction = path.googleModel(sub);
+    if (!modelAction) {
+      logger.debug(`Non-model Google path, cannot route to CCA: ${sub}`);
+      return denied("Google (unsupported path)");
+    }
+
+    const unwrapThenRewrite = withUnwrap(rewrite);
+    const orderedStrategies = getOrderedStrategies(account, modelAction.model);
+
+    if (orderedStrategies.length === 0) {
+      const now = Date.now();
+      let minWait = COOLDOWN_MS;
+      for (const s of strategies) {
+        const until = cooldowns.get(cooldownKey(account, s));
+        if (until) minWait = Math.min(minWait, Math.max(0, until - now));
+      }
+      const retryAfterS = Math.ceil(minWait / 1000);
+      return new Response(
+        JSON.stringify({
+          error: { message: "All Google strategies cooling down", type: "rate_limit_error", code: "429" },
+        }),
+        { status: 429, headers: { "Content-Type": "application/json", "Retry-After": String(retryAfterS) } },
+      );
+    }
+
+    let saw429 = false;
+    let lastResponse: Response | null = null;
+
+    for (const strategy of orderedStrategies) {
+      const model = strategy.modelMapper ? strategy.modelMapper(modelAction.model) : modelAction.model;
+      const isImageModel = model.includes("image");
+      const wrapOpts = isImageModel ? { ...strategy.wrapOpts, requestType: "image_gen" as const } : strategy.wrapOpts;
+      const requestBody = maybeWrap(body.parsed, body.forwardBody, projectId, model, wrapOpts);
+
+      const headers: Record<string, string> = {
+        ...strategy.headers,
+        Authorization: `Bearer ${accessToken}`,
+        "Content-Type": "application/json",
+        Accept: body.stream ? "text/event-stream" : "application/json",
+      };
+
+      logger.info(`Google strategy=${strategy.name} account=${account} model=${model}`);
+
+      // Antigravity forces streaming endpoint for gemini-3-pro* and image models (matches CLIProxyAPI behavior)
+      const forceStream = strategy.name === "antigravity" && (model.startsWith("gemini-3-pro") || isImageModel);
+      const action = forceStream ? "streamGenerateContent" : modelAction.action;
+
+      for (const endpoint of strategy.endpoints) {
+        const url = buildUrl(endpoint, action);
+        try {
+          const forceStreamNonStreaming = forceStream && !body.stream;
+          const response = await forward({
+            url,
+            body: requestBody,
+            streaming: forceStreamNonStreaming ? true : body.stream,
+            headers,
+            providerName: `Google/${strategy.name}`,
+            rewrite: unwrapThenRewrite,
+            email,
+          });
+
+          // When we forced streaming but client expects JSON, buffer SSE and return last chunk
+          if (forceStreamNonStreaming && response.ok && response.body) {
+            const merged = await bufferSSEToJSON(response);
+            if (merged) {
+              markSuccess(account, strategy);
+              return merged;
+            }
+          }
+
+          if (response.status === 401 || response.status === 403) {
+            return response;
+          }
+
+          if (response.status === 404) {
+            lastResponse = response;
+            logger.debug(`Google strategy=${strategy.name} model not found (404), trying next endpoint`);
+            continue;
+          }
+
+          if (response.status === 429) {
+            saw429 = true;
+            lastResponse = response;
+            logger.debug(`Google strategy=${strategy.name} failed (${response.status}), trying next`);
+            break;
+          }
+
+          if (response.status >= 500) {
+            lastResponse = response;
+            logger.debug(`Google strategy=${strategy.name} failed (${response.status}), trying next`);
+            continue;
+          }
+
+          markSuccess(account, strategy);
+          return response;
+        } catch (err) {
+          lastResponse = apiError(502, `Google/${strategy.name} endpoint failed: ${String(err)}`);
+          logger.debug(`Google strategy=${strategy.name} endpoint error`, { error: String(err) });
+        }
+      }
+
+      markFailure(account, strategy);
+    }
+
+    if (saw429) {
+      return lastResponse ?? apiError(429, "All Google strategies rate limited", "rate_limit_error");
+    }
+
+    return lastResponse ?? apiError(502, "All Google strategies exhausted");
+  },
+};

package/src/routing/affinity.ts

@@ -8,7 +8,7 @@ import type { QuotaPool } from "./cooldown.ts";
 interface AffinityEntry {
   pool: QuotaPool;
   account: number;
-  assignedAt: number;
+  lastUsedAt: number;
 }
 
 /** Affinity expires after 2 hours of inactivity. */
@@ -51,7 +51,7 @@ class AffinityStore {
     const k = this.key(threadId, ampProvider);
     const entry = this.map.get(k);
     if (!entry) return undefined;
-    if (Date.now() - entry.assignedAt > TTL_MS) {
+    if (Date.now() - entry.lastUsedAt > TTL_MS) {
       this.removeExpired(k, entry);
       return undefined;
     }
@@ -61,7 +61,7 @@ class AffinityStore {
   /** Read affinity and touch (extend TTL). */
   get(threadId: string, ampProvider: string): AffinityEntry | undefined {
     const entry = this.peek(threadId, ampProvider);
-    if (entry) entry.assignedAt = Date.now();
+    if (entry) entry.lastUsedAt = Date.now();
     return entry;
   }
 
@@ -76,7 +76,7 @@ class AffinityStore {
     } else {
       this.incCount(pool, account);
     }
-    this.map.set(k, { pool, account, assignedAt: Date.now() });
+    this.map.set(k, { pool, account, lastUsedAt: Date.now() });
   }
 
   /** Break affinity when account is exhausted — allow re-routing. */
@@ -100,7 +100,7 @@ class AffinityStore {
     this.cleanupTimer = setInterval(() => {
       const now = Date.now();
       for (const [k, entry] of this.map) {
-        if (now - entry.assignedAt > TTL_MS) {
+        if (now - entry.lastUsedAt > TTL_MS) {
           this.removeExpired(k, entry);
         }
       }

package/src/routing/cooldown.ts

@@ -3,7 +3,7 @@
 
 import { logger } from "../utils/logger.ts";
 
-export type QuotaPool = "anthropic" | "codex" | "gemini" | "antigravity";
+export type QuotaPool = "anthropic" | "codex" | "google";
 
 interface CooldownEntry {
   until: number;
@@ -74,6 +74,22 @@ class CooldownTracker {
     logger.warn(`Account disabled (403): ${k}`, { cooldownHours: FORBIDDEN_COOLDOWN_MS / 3600_000 });
   }
 
+  /** Return shortest remaining wait (ms) among non-exhausted entries for the given candidates.
+   *  Returns undefined if no candidates are cooling down or all are exhausted. */
+  shortestBurstWait(candidates: { pool: QuotaPool; account: number }[]): number | undefined {
+    let shortest: number | undefined;
+    const now = Date.now();
+    for (const c of candidates) {
+      const entry = this.entries.get(this.key(c.pool, c.account));
+      if (!entry || entry.exhausted) continue;
+      const remaining = entry.until - now;
+      if (remaining > 0 && (shortest === undefined || remaining < shortest)) {
+        shortest = remaining;
+      }
+    }
+    return shortest;
+  }
+
   recordSuccess(pool: QuotaPool, account: number): void {
     this.entries.delete(this.key(pool, account));
   }

package/src/routing/retry.ts

@@ -4,12 +4,14 @@ import type { ProxyConfig } from "../config/config.ts";
 import type { ParsedBody } from "../server/body.ts";
 import { logger } from "../utils/logger.ts";
 import { cooldown, parseRetryAfter, type QuotaPool } from "./cooldown.ts";
-import { type RouteResult, recordSuccess, reroute } from "./router.ts";
+import { buildCandidates, type RouteResult, recordSuccess, reroute } from "./router.ts";
 
 /** Max reroute attempts before falling back to upstream. */
 const MAX_REROUTE_ATTEMPTS = 4;
 /** Max seconds to wait-and-retry on the same account (preserves prompt cache). */
 const CACHE_PRESERVE_WAIT_MAX_S = 10;
+/** Max ms to wait when all candidates are burst-cooling before giving up. */
+const BURST_WAIT_MAX_MS = 30_000;
 
 /** Status codes that trigger rerouting to a different account/pool. */
 const REROUTABLE_STATUSES = new Set([429, 403]);
@@ -41,7 +43,7 @@ export async function tryWithCachePreserve(
   await Bun.sleep(retryAfter * 1000);
   const response = await route.handler!.forward(sub, body, headers, rewrite, route.account);
 
-  if (response.status !== 429 && response.status !== 401) {
+  if (response.status !== 429 && response.status !== 403 && response.status !== 401) {
     recordSuccess(route.pool!, route.account);
     return response;
   }
@@ -49,6 +51,9 @@ export async function tryWithCachePreserve(
     const nextRetryAfter = parseRetryAfter(response.headers.get("retry-after"));
     cooldown.record429(route.pool!, route.account, nextRetryAfter);
   }
+  if (response.status === 403) {
+    cooldown.record403(route.pool!, route.account);
+  }
   return null;
 }
 
@@ -64,7 +69,19 @@ export async function tryReroute(
   let currentAccount = initialRoute.account;
 
   for (let attempt = 0; attempt < MAX_REROUTE_ATTEMPTS; attempt++) {
-    const next = reroute(ctx.providerName, ctx.ampModel, ctx.config, currentPool, currentAccount, ctx.threadId);
+    let next = reroute(ctx.providerName, ctx.ampModel, ctx.config, currentPool, currentAccount, ctx.threadId);
+
+    // All candidates cooling down — wait for shortest burst then retry
+    if (!next?.handler) {
+      const candidates = buildCandidates(ctx.providerName, ctx.config);
+      const waitMs = cooldown.shortestBurstWait(candidates);
+      if (waitMs && waitMs <= BURST_WAIT_MAX_MS) {
+        logger.info(`All accounts cooling, waiting ${Math.ceil(waitMs / 1000)}s for burst cooldown`);
+        await Bun.sleep(waitMs + 100); // small buffer
+        next = reroute(ctx.providerName, ctx.ampModel, ctx.config, currentPool, currentAccount, ctx.threadId);
+      }
+    }
+
     if (!next?.handler) break;
 
     logger.info(`REROUTE (${status}) -> ${next.decision} account=${next.account}`);

package/src/routing/router.ts

@@ -5,17 +5,16 @@
  * - Thread affinity: same thread sticks to same account
  * - Cooldown: skip 429'd accounts, detect quota exhaustion
  * - Least-connections: prefer accounts with fewer active threads
- * - Google cascade: gemini accounts → antigravity accounts (separate quotas)
+ * - Google: single pool with internal strategy fallback (gemini/antigravity)
  */
 
 import type { ProviderName } from "../auth/store.ts";
 import * as store from "../auth/store.ts";
 import type { ProxyConfig } from "../config/config.ts";
 import { provider as anthropic } from "../providers/anthropic.ts";
-import { provider as antigravity } from "../providers/antigravity.ts";
 import type { Provider } from "../providers/base.ts";
 import { provider as codex } from "../providers/codex.ts";
-import { provider as gemini } from "../providers/gemini.ts";
+import { provider as google } from "../providers/google.ts";
 import { logger, type RouteDecision } from "../utils/logger.ts";
 import { affinity } from "./affinity.ts";
 import { cooldown, type QuotaPool } from "./cooldown.ts";
@@ -46,10 +45,7 @@ const PROVIDER_REGISTRY = new Map<string, { configKey: keyof ProxyConfig["provid
     "google",
     {
       configKey: "google",
-      entries: [
-        { provider: gemini, pool: "gemini", credentialName: "google" },
-        { provider: antigravity, pool: "antigravity", credentialName: "google" },
-      ],
+      entries: [{ provider: google, pool: "google", credentialName: "google" }],
     },
   ],
 ]);
@@ -85,6 +81,13 @@ export function routeRequest(
 ): RouteResult {
   const modelStr = model ?? "unknown";
 
+  // Early exit if provider is disabled in config
+  const reg = PROVIDER_REGISTRY.get(ampProvider);
+  if (!reg || !config.providers[reg.configKey]) {
+    logger.route("AMP_UPSTREAM", ampProvider, modelStr);
+    return result(null, ampProvider, modelStr, 0, null);
+  }
+
   // Check thread affinity (keyed by threadId + ampProvider)
   if (threadId) {
     const pinned = affinity.get(threadId, ampProvider);
@@ -153,7 +156,7 @@ export function recordSuccess(pool: QuotaPool, account: number): void {
   cooldown.recordSuccess(pool, account);
 }
 
-function buildCandidates(ampProvider: string, config: ProxyConfig): Candidate[] {
+export function buildCandidates(ampProvider: string, config: ProxyConfig): Candidate[] {
   const reg = PROVIDER_REGISTRY.get(ampProvider);
   if (!reg || !config.providers[reg.configKey]) return [];
 
@@ -183,18 +186,21 @@ function pickCandidate(candidates: Candidate[]): Candidate | null {
   if (available.length === 0) return null;
 
   // Pick the one with least active threads (least-connections)
-  let best = available[0]!;
-  let bestLoad = affinity.activeCount(best.pool, best.account);
+  // When multiple candidates have the same load, pick randomly for even distribution
+  let bestLoad = affinity.activeCount(available[0]!.pool, available[0]!.account);
+  let ties: Candidate[] = [available[0]!];
 
   for (let i = 1; i < available.length; i++) {
     const load = affinity.activeCount(available[i]!.pool, available[i]!.account);
     if (load < bestLoad) {
-      best = available[i]!;
       bestLoad = load;
+      ties = [available[i]!];
+    } else if (load === bestLoad) {
+      ties.push(available[i]!);
     }
   }
 
-  return best;
+  return ties[Math.floor(Math.random() * ties.length)]!;
 }
 
 function providerForPool(pool: QuotaPool): Provider | null {

package/src/server/body.ts

@@ -3,7 +3,7 @@
  *  Slow path: full JSON.parse only when .parsed or .forwardBody is accessed
  *  (e.g. Google CCA wrapping, model rewrite). */
 
-import { resolveModel, rewriteBodyModel } from "../routing/models.ts";
+import { resolveModel, rewriteBodyModel } from "../utils/models.ts";
 import * as path from "../utils/path.ts";
 
 export interface ParsedBody {

package/src/server/server.ts

@@ -10,6 +10,7 @@ import { recordSuccess, routeRequest } from "../routing/router.ts";
 import { handleInternal, isLocalMethod } from "../tools/internal.ts";
 import { logger } from "../utils/logger.ts";
 import * as path from "../utils/path.ts";
+import { apiError } from "../utils/responses.ts";
 import { stats } from "../utils/stats.ts";
 import { type ParsedBody, parseBody } from "./body.ts";
 
@@ -29,7 +30,7 @@ export function startServer(config: ProxyConfig): ReturnType<typeof Bun.serve> {
         return response;
       } catch (err) {
         logger.error("Unhandled server error", { error: String(err) });
-        return Response.json({ error: "Internal proxy error" }, { status });
+        return apiError(status, "Internal proxy error");
       } finally {
         logger.info(`${req.method} ${url.pathname}${url.search} ${status}`, { duration: Date.now() - startTime });
       }
@@ -101,7 +102,10 @@ async function handleProvider(
     const rewrite = ampModel ? rewriter.rewrite(ampModel) : undefined;
     const handlerResponse = await route.handler.forward(sub, body, req.headers, rewrite, route.account);
 
-    if ((handlerResponse.status === 429 || handlerResponse.status === 403) && route.pool) {
+    if (
+      (handlerResponse.status === 429 || handlerResponse.status === 403 || handlerResponse.status === 404) &&
+      route.pool
+    ) {
       const ctx = { providerName, ampModel, config, sub, body, headers: req.headers, rewrite, threadId };
       // 429: try short wait to preserve prompt cache first
       const cached =

package/src/utils/code-assist.ts

@@ -6,18 +6,23 @@ interface WrapOptions {
   body: Record<string, unknown>;
   userAgent: "antigravity" | "pi-coding-agent";
   requestIdPrefix: "agent" | "pi";
-  requestType?: "agent";
+  requestType?: "agent" | "image_gen";
 }
 
 /** Wrap a raw request body in the Cloud Code Assist envelope. */
 function wrapRequest(opts: WrapOptions): string {
+  const isImageGen = opts.requestType === "image_gen";
+  const requestId = isImageGen
+    ? `image_gen/${Date.now()}/${crypto.randomUUID()}/12`
+    : `${opts.requestIdPrefix}-${Date.now()}-${crypto.randomUUID().slice(0, 8)}`;
+
   return JSON.stringify({
     project: opts.projectId,
     model: opts.model,
     request: opts.body,
     ...(opts.requestType && { requestType: opts.requestType }),
     userAgent: opts.userAgent,
-    requestId: `${opts.requestIdPrefix}-${Date.now()}-${crypto.randomUUID().slice(0, 8)}`,
+    requestId,
   });
 }
 
@@ -118,7 +123,11 @@ export function maybeWrap(
   raw: string,
   projectId: string,
   model: string,
-  opts: { userAgent: "antigravity" | "pi-coding-agent"; requestIdPrefix: "agent" | "pi"; requestType?: "agent" },
+  opts: {
+    userAgent: "antigravity" | "pi-coding-agent";
+    requestIdPrefix: "agent" | "pi";
+    requestType?: "agent" | "image_gen";
+  },
 ): string {
   if (!parsed) return raw;
   if (parsed.project) return raw;

package/src/utils/logger.ts

@@ -1,6 +1,6 @@
 /** Structured logging with route decision tracking. */
 
-export type RouteDecision = "LOCAL_CLAUDE" | "LOCAL_CODEX" | "LOCAL_GEMINI" | "LOCAL_ANTIGRAVITY" | "AMP_UPSTREAM";
+export type RouteDecision = "LOCAL_CLAUDE" | "LOCAL_CODEX" | "LOCAL_GOOGLE" | "AMP_UPSTREAM";
 
 export type LogLevel = "debug" | "info" | "warn" | "error";
 
@@ -41,8 +41,7 @@ const LEVEL_COLORS: Record<LogLevel, string> = {
 const ROUTE_COLORS: Record<RouteDecision, string> = {
   LOCAL_CLAUDE: GREEN,
   LOCAL_CODEX: GREEN,
-  LOCAL_GEMINI: GREEN,
-  LOCAL_ANTIGRAVITY: GREEN,
+  LOCAL_GOOGLE: GREEN,
   AMP_UPSTREAM: YELLOW,
 };
 

package/src/utils/path.ts

@@ -5,7 +5,7 @@ import { browserPrefixes, passthroughExact, passthroughPrefixes } from "../const
 const PROVIDER_RE = /^\/api\/provider\/([^/]+)/;
 const SUBPATH_RE = /^\/api\/provider\/[^/]+(\/.*)/;
 const MODEL_RE = /models\/([^/:]+)/;
-const GEMINI_RE = /models\/([^/:]+):(\w+)/;
+const GOOGLE_MODEL_RE = /models\/([^/:]+):(\w+)/;
 
 export function passthrough(pathname: string): boolean {
   if ((passthroughExact as readonly string[]).includes(pathname)) return true;
@@ -13,8 +13,7 @@ export function passthrough(pathname: string): boolean {
 }
 
 export function browser(pathname: string): boolean {
-  if ((passthroughExact as readonly string[]).includes(pathname)) return true;
-  return browserPrefixes.some((prefix) => pathname.startsWith(prefix));
+  return browserPrefixes.some((prefix) => pathname === prefix || pathname.startsWith(`${prefix}/`));
 }
 
 export function provider(pathname: string): string | null {
@@ -32,8 +31,8 @@ export function modelFromUrl(url: string): string | null {
   return match?.[1] ?? null;
 }
 
-export function gemini(url: string): { model: string; action: string } | null {
-  const match = url.match(GEMINI_RE);
+export function googleModel(url: string): { model: string; action: string } | null {
+  const match = url.match(GOOGLE_MODEL_RE);
   if (!match) return null;
   return { model: match[1]!, action: match[2]! };
 }

package/src/utils/responses.ts

@@ -0,0 +1,3 @@
+export function apiError(status: number, message: string, type = "api_error"): Response {
+  return Response.json({ error: { message, type, code: String(status) } }, { status });
+}

package/src/providers/antigravity.ts

@@ -1,87 +0,0 @@
-/** Forwards requests to Cloud Code Assist API using Antigravity quota.
- *  Uses the shared Google OAuth token with Antigravity headers/endpoints.
- *  Tries multiple endpoints with fallback on 5xx. */
-
-import { google as config } from "../auth/configs.ts";
-import * as oauth from "../auth/oauth.ts";
-import * as store from "../auth/store.ts";
-import { ANTIGRAVITY_DAILY_ENDPOINT, AUTOPUSH_ENDPOINT, CODE_ASSIST_ENDPOINT } from "../constants.ts";
-import { buildUrl, maybeWrap, withUnwrap } from "../utils/code-assist.ts";
-import { logger } from "../utils/logger.ts";
-import * as path from "../utils/path.ts";
-import type { Provider } from "./base.ts";
-import { denied, forward } from "./forward.ts";
-
-const endpoints = [ANTIGRAVITY_DAILY_ENDPOINT, AUTOPUSH_ENDPOINT, CODE_ASSIST_ENDPOINT];
-
-const antigravityHeaders: Readonly<Record<string, string>> = {
-  "User-Agent": "antigravity/1.15.8 darwin/arm64",
-  "X-Goog-Api-Client": "google-cloud-sdk vscode_cloudshelleditor/0.1",
-  "Client-Metadata": JSON.stringify({
-    ideType: "IDE_UNSPECIFIED",
-    platform: "PLATFORM_UNSPECIFIED",
-    pluginType: "GEMINI",
-  }),
-};
-
-export const provider: Provider = {
-  name: "Antigravity",
-  routeDecision: "LOCAL_ANTIGRAVITY",
-
-  isAvailable: (account?: number) =>
-    account !== undefined ? !!store.get("google", account)?.refreshToken : oauth.ready(config),
-
-  accountCount: () => oauth.accountCount(config),
-
-  async forward(sub, body, _originalHeaders, rewrite, account = 0) {
-    const accessToken = await oauth.token(config, account);
-    if (!accessToken) return denied("Antigravity");
-
-    const creds = store.get("google", account);
-    const projectId = creds?.projectId ?? "";
-
-    const headers: Record<string, string> = {
-      ...antigravityHeaders,
-      Authorization: `Bearer ${accessToken}`,
-      "Content-Type": "application/json",
-      Accept: body.stream ? "text/event-stream" : "application/json",
-    };
-    const gemini = path.gemini(sub);
-    const action = gemini?.action ?? "generateContent";
-    const model = gemini?.model === "gemini-3-flash-preview" ? "gemini-3-flash" : (gemini?.model ?? "");
-    const requestBody = maybeWrap(body.parsed, body.forwardBody, projectId, model, {
-      userAgent: "antigravity",
-      requestIdPrefix: "agent",
-      requestType: "agent",
-    });
-    const unwrapThenRewrite = withUnwrap(rewrite);
-
-    return tryEndpoints(requestBody, body.stream, headers, action, unwrapThenRewrite, creds?.email);
-  },
-};
-
-async function tryEndpoints(
-  body: string,
-  streaming: boolean,
-  headers: Record<string, string>,
-  action: string,
-  rewrite?: (data: string) => string,
-  email?: string,
-): Promise<Response> {
-  let lastError: Error | null = null;
-
-  for (const endpoint of endpoints) {
-    const url = buildUrl(endpoint, action);
-    try {
-      const response = await forward({ url, body, streaming, headers, providerName: "Antigravity", rewrite, email });
-      if (response.status < 500) return response;
-      lastError = new Error(`${endpoint} returned ${response.status}`);
-      logger.debug("Endpoint 5xx, trying next", { provider: "Antigravity" });
-    } catch (err) {
-      lastError = err instanceof Error ? err : new Error(String(err));
-      logger.debug("Endpoint failed, trying next", { error: String(err) });
-    }
-  }
-
-  return Response.json({ error: `All Antigravity endpoints failed: ${lastError?.message}` }, { status: 502 });
-}

package/src/providers/gemini.ts

@@ -1,73 +0,0 @@
-/** Forwards requests to Cloud Code Assist API using Gemini CLI quota.
- *  Amp CLI uses @google/genai SDK with vertexai:true — sends Vertex AI format
- *  (e.g. /v1beta1/publishers/google/models/{model}:streamGenerateContent?alt=sse).
- *  We wrap the native body in a CCA envelope, forward to cloudcode-pa, and unwrap
- *  the response so Amp CLI sees standard Vertex AI SSE chunks. */
-
-import { google as config } from "../auth/configs.ts";
-import * as oauth from "../auth/oauth.ts";
-import * as store from "../auth/store.ts";
-import { CODE_ASSIST_ENDPOINT } from "../constants.ts";
-import { buildUrl, maybeWrap, withUnwrap } from "../utils/code-assist.ts";
-import { logger } from "../utils/logger.ts";
-import * as path from "../utils/path.ts";
-import type { Provider } from "./base.ts";
-import { denied, forward } from "./forward.ts";
-
-const geminiHeaders: Readonly<Record<string, string>> = {
-  "User-Agent": "google-cloud-sdk vscode_cloudshelleditor/0.1",
-  "X-Goog-Api-Client": "gl-node/22.17.0",
-  "Client-Metadata": JSON.stringify({
-    ideType: "IDE_UNSPECIFIED",
-    platform: "PLATFORM_UNSPECIFIED",
-    pluginType: "GEMINI",
-  }),
-};
-
-export const provider: Provider = {
-  name: "Gemini CLI",
-  routeDecision: "LOCAL_GEMINI",
-
-  isAvailable: (account?: number) =>
-    account !== undefined ? !!store.get("google", account)?.refreshToken : oauth.ready(config),
-
-  accountCount: () => oauth.accountCount(config),
-
-  async forward(sub, body, _originalHeaders, rewrite, account = 0) {
-    const accessToken = await oauth.token(config, account);
-    if (!accessToken) return denied("Gemini CLI");
-
-    const creds = store.get("google", account);
-    const projectId = creds?.projectId ?? "";
-
-    const headers: Record<string, string> = {
-      ...geminiHeaders,
-      Authorization: `Bearer ${accessToken}`,
-      "Content-Type": "application/json",
-      Accept: body.stream ? "text/event-stream" : "application/json",
-    };
-
-    const gemini = path.gemini(sub);
-    if (!gemini) {
-      logger.debug(`Non-model Gemini path, cannot route to CCA: ${sub}`);
-      return denied("Gemini CLI (unsupported path)");
-    }
-
-    const url = buildUrl(CODE_ASSIST_ENDPOINT, gemini.action);
-    const requestBody = maybeWrap(body.parsed, body.forwardBody, projectId, gemini.model, {
-      userAgent: "pi-coding-agent",
-      requestIdPrefix: "pi",
-    });
-    const unwrapThenRewrite = withUnwrap(rewrite);
-
-    return forward({
-      url,
-      body: requestBody,
-      streaming: body.stream,
-      headers,
-      providerName: "Gemini CLI",
-      rewrite: unwrapThenRewrite,
-      email: store.get("google", account)?.email,
-    });
-  },
-};