ampcode-connector 0.1.0

package/src/providers/antigravity.ts

@@ -0,0 +1,110 @@
+/** 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 * as codeAssist 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 "./base.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: "text/event-stream",
+    };
+
+    const anthropicBeta = originalHeaders.get("anthropic-beta");
+    if (anthropicBeta) headers["anthropic-beta"] = anthropicBeta;
+
+    const gemini = path.gemini(sub);
+    const action = gemini?.action ?? "generateContent";
+    const model = gemini?.model ?? "";
+    const requestBody = maybeWrap(body, projectId, model);
+    const unwrapThenRewrite = withUnwrap(rewrite);
+
+    return tryEndpoints(requestBody, headers, action, unwrapThenRewrite);
+  },
+};
+
+function withUnwrap(rewrite?: (d: string) => string): (d: string) => string {
+  return rewrite ? (d: string) => rewrite(codeAssist.unwrap(d)) : codeAssist.unwrap;
+}
+
+function maybeWrap(body: string, projectId: string, model: string): string {
+  try {
+    const parsed = JSON.parse(body) as Record<string, unknown>;
+    if (parsed["project"]) return body;
+    return codeAssist.wrapRequest({
+      projectId,
+      model,
+      body: parsed,
+      userAgent: "antigravity",
+      requestIdPrefix: "agent",
+      requestType: "agent",
+    });
+  } catch (err) {
+    logger.debug("Body parse failed, forwarding as-is", { error: String(err) });
+    return body;
+  }
+}
+
+async function tryEndpoints(
+  body: string,
+  headers: Record<string, string>,
+  action: string,
+  rewrite?: (data: string) => string,
+): Promise<Response> {
+  let lastError: Error | null = null;
+
+  for (const endpoint of endpoints) {
+    const url = codeAssist.buildUrl(endpoint, action);
+    try {
+      const response = await forward({ url, body, headers, providerName: "Antigravity", rewrite });
+      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 new Response(JSON.stringify({ error: `All Antigravity endpoints failed: ${lastError?.message}` }), {
+    status: 502,
+    headers: { "Content-Type": "application/json" },
+  });
+}

package/src/providers/base.ts

@@ -0,0 +1,68 @@
+/** Provider interface and shared request forwarding. */
+
+import type { RouteDecision } from "../utils/logger.ts";
+import { logger } from "../utils/logger.ts";
+import * as path from "../utils/path.ts";
+import * as sse from "../utils/streaming.ts";
+
+export interface Provider {
+  readonly name: string;
+  readonly routeDecision: RouteDecision;
+  isAvailable(account?: number): boolean;
+  accountCount(): number;
+  forward(
+    path: string,
+    body: string,
+    headers: Headers,
+    rewrite?: (data: string) => string,
+    account?: number,
+  ): Promise<Response>;
+}
+
+interface ForwardOptions {
+  url: string;
+  body: string;
+  headers: Record<string, string>;
+  providerName: string;
+  rewrite?: (data: string) => string;
+}
+
+export async function forward(opts: ForwardOptions): Promise<Response> {
+  const response = await fetch(opts.url, {
+    method: "POST",
+    headers: opts.headers,
+    body: opts.body,
+  });
+
+  const isSSE = response.headers.get("Content-Type")?.includes("text/event-stream") || path.streaming(opts.body);
+  if (isSSE) return sse.proxy(response, opts.rewrite);
+
+  if (!response.ok) {
+    const text = await response.text();
+    logger.error(`${opts.providerName} API error`, { error: text.slice(0, 200) });
+    return new Response(text, {
+      status: response.status,
+      headers: { "Content-Type": response.headers.get("Content-Type") ?? "application/json" },
+    });
+  }
+
+  if (opts.rewrite) {
+    const text = await response.text();
+    return new Response(opts.rewrite(text), {
+      status: response.status,
+      headers: { "Content-Type": response.headers.get("Content-Type") ?? "application/json" },
+    });
+  }
+
+  return new Response(response.body, {
+    status: response.status,
+    headers: { "Content-Type": response.headers.get("Content-Type") ?? "application/json" },
+  });
+}
+
+export function denied(providerName: string): Response {
+  return new Response(JSON.stringify({ error: `No ${providerName} OAuth token available. Run login first.` }), {
+    status: 401,
+    headers: { "Content-Type": "application/json" },
+  });
+}

package/src/providers/codex.ts

@@ -0,0 +1,36 @@
+/** Forwards requests to api.openai.com with Codex CLI OAuth token. */
+
+import { codex as config } from "../auth/configs.ts";
+import * as oauth from "../auth/oauth.ts";
+import * as store from "../auth/store.ts";
+import { OPENAI_API_URL } from "../constants.ts";
+import * as path from "../utils/path.ts";
+import type { Provider } from "./base.ts";
+import { denied, forward } from "./base.ts";
+
+export const provider: Provider = {
+  name: "OpenAI Codex",
+  routeDecision: "LOCAL_CODEX",
+
+  isAvailable: (account?: number) =>
+    account !== undefined ? !!store.get("codex", 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("OpenAI Codex");
+
+    return forward({
+      url: `${OPENAI_API_URL}${sub}`,
+      body,
+      providerName: "OpenAI Codex",
+      rewrite,
+      headers: {
+        "Content-Type": "application/json",
+        Authorization: `Bearer ${accessToken}`,
+        Accept: path.streaming(body) ? "text/event-stream" : "application/json",
+      },
+    });
+  },
+};

package/src/providers/gemini.ts

@@ -0,0 +1,83 @@
+/** 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 * as codeAssist 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 "./base.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: "text/event-stream",
+    };
+
+    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 = codeAssist.buildUrl(CODE_ASSIST_ENDPOINT, gemini.action);
+    const requestBody = maybeWrap(body, projectId, gemini.model);
+    const unwrapThenRewrite = withUnwrap(rewrite);
+
+    return forward({ url, body: requestBody, headers, providerName: "Gemini CLI", rewrite: unwrapThenRewrite });
+  },
+};
+
+function withUnwrap(rewrite?: (d: string) => string): (d: string) => string {
+  return rewrite ? (d: string) => rewrite(codeAssist.unwrap(d)) : codeAssist.unwrap;
+}
+
+function maybeWrap(body: string, projectId: string, model: string): string {
+  try {
+    const parsed = JSON.parse(body) as Record<string, unknown>;
+    if (parsed["project"]) return body;
+    return codeAssist.wrapRequest({
+      projectId,
+      model,
+      body: parsed,
+      userAgent: "pi-coding-agent",
+      requestIdPrefix: "pi",
+    });
+  } catch (err) {
+    logger.debug("Body parse failed, forwarding as-is", { error: String(err) });
+    return body;
+  }
+}

package/src/proxy/rewriter.ts

@@ -0,0 +1,69 @@
+/**
+ * SSE response rewriting: model name substitution + thinking block suppression.
+ *
+ * Thinking blocks are filtered when tool_use is present because the Amp client
+ * struggles with both simultaneously (ref: CLIProxyAPI response_rewriter.go:72-94).
+ */
+
+import { modelFieldPaths } from "../constants.ts";
+
+export function rewrite(originalModel: string): (data: string) => string {
+  return (data: string) => {
+    if (data === "[DONE]") return data;
+
+    try {
+      const parsed = JSON.parse(data) as Record<string, unknown>;
+      let modified = false;
+
+      for (const path of modelFieldPaths) {
+        const current = getField(parsed, path);
+        if (current !== undefined && current !== originalModel) {
+          setField(parsed, path, originalModel);
+          modified = true;
+        }
+      }
+
+      if (suppressThinking(parsed)) modified = true;
+
+      return modified ? JSON.stringify(parsed) : data;
+    } catch {
+      return data;
+    }
+  };
+}
+
+function getField(obj: Record<string, unknown>, path: string): unknown {
+  const parts = path.split(".");
+  let current: unknown = obj;
+  for (const part of parts) {
+    if (current == null || typeof current !== "object") return undefined;
+    current = (current as Record<string, unknown>)[part];
+  }
+  return current;
+}
+
+function setField(obj: Record<string, unknown>, path: string, value: unknown): void {
+  const parts = path.split(".");
+  let current: unknown = obj;
+  for (let i = 0; i < parts.length - 1; i++) {
+    if (current == null || typeof current !== "object") return;
+    current = (current as Record<string, unknown>)[parts[i]!];
+  }
+  if (current != null && typeof current === "object") {
+    (current as Record<string, unknown>)[parts[parts.length - 1]!] = value;
+  }
+}
+
+function suppressThinking(data: Record<string, unknown>): boolean {
+  const content = data["content"];
+  if (!Array.isArray(content)) return false;
+
+  const hasToolUse = content.some((b: Record<string, unknown>) => b["type"] === "tool_use");
+  if (!hasToolUse) return false;
+
+  const hasThinking = content.some((b: Record<string, unknown>) => b["type"] === "thinking");
+  if (!hasThinking) return false;
+
+  data["content"] = content.filter((b: Record<string, unknown>) => b["type"] !== "thinking");
+  return true;
+}

package/src/proxy/upstream.ts

@@ -0,0 +1,41 @@
+/** Reverse proxy to ampcode.com for non-intercepted routes and fallback. */
+
+import { logger } from "../utils/logger.ts";
+
+export async function forward(request: Request, ampUpstreamUrl: string, ampApiKey?: string): Promise<Response> {
+  const url = new URL(request.url);
+  const upstreamUrl = new URL(url.pathname + url.search, ampUpstreamUrl);
+
+  const upstreamHost = new URL(ampUpstreamUrl).host;
+  const headers = new Headers(request.headers);
+  if (ampApiKey) headers.set("Authorization", `Bearer ${ampApiKey}`);
+  headers.set("Host", upstreamHost);
+
+  logger.debug("Forwarding to Amp upstream", { provider: "amp" });
+
+  try {
+    const response = await fetch(upstreamUrl.toString(), {
+      method: request.method,
+      headers,
+      redirect: "manual",
+      body: request.method !== "GET" && request.method !== "HEAD" ? request.body : undefined,
+      duplex: "half" as const,
+    });
+
+    const responseHeaders = new Headers(response.headers);
+    responseHeaders.delete("Content-Encoding");
+    responseHeaders.delete("Content-Length");
+
+    return new Response(response.body, {
+      status: response.status,
+      statusText: response.statusText,
+      headers: responseHeaders,
+    });
+  } catch (err) {
+    logger.error("Upstream proxy error", { error: String(err) });
+    return new Response(JSON.stringify({ error: "Failed to connect to Amp upstream", details: String(err) }), {
+      status: 502,
+      headers: { "Content-Type": "application/json" },
+    });
+  }
+}

package/src/routing/affinity.ts

@@ -0,0 +1,56 @@
+/** Thread → (quotaPool, account) affinity map.
+ *  Ensures a thread sticks to the same account for session consistency. */
+
+import type { QuotaPool } from "./cooldown.ts";
+
+interface AffinityEntry {
+  pool: QuotaPool;
+  account: number;
+  assignedAt: number;
+}
+
+/** Affinity expires after 2 hours of inactivity. */
+const TTL_MS = 2 * 3600_000;
+/** Cleanup stale entries every 10 minutes. */
+const CLEANUP_INTERVAL_MS = 10 * 60_000;
+
+const map = new Map<string, AffinityEntry>();
+
+export function get(threadId: string): AffinityEntry | undefined {
+  const entry = map.get(threadId);
+  if (!entry) return undefined;
+  if (Date.now() - entry.assignedAt > TTL_MS) {
+    map.delete(threadId);
+    return undefined;
+  }
+  return entry;
+}
+
+export function set(threadId: string, pool: QuotaPool, account: number): void {
+  map.set(threadId, { pool, account, assignedAt: Date.now() });
+}
+
+/** Break affinity when account is exhausted — allow re-routing. */
+export function clear(threadId: string): void {
+  map.delete(threadId);
+}
+
+/** Count active threads pinned to a specific (pool, account). */
+export function activeCount(pool: QuotaPool, account: number): number {
+  const now = Date.now();
+  let count = 0;
+  for (const entry of map.values()) {
+    if (entry.pool === pool && entry.account === account && now - entry.assignedAt <= TTL_MS) {
+      count++;
+    }
+  }
+  return count;
+}
+
+/** Periodic cleanup of expired entries. */
+setInterval(() => {
+  const now = Date.now();
+  for (const [threadId, entry] of map) {
+    if (now - entry.assignedAt > TTL_MS) map.delete(threadId);
+  }
+}, CLEANUP_INTERVAL_MS);

package/src/routing/cooldown.ts

@@ -0,0 +1,88 @@
+/** Per-(quotaPool, account) cooldown tracking.
+ *  Distinguishes short burst 429s from quota exhaustion. */
+
+import { logger } from "../utils/logger.ts";
+
+export type QuotaPool = "anthropic" | "codex" | "gemini" | "antigravity";
+
+interface CooldownEntry {
+  until: number;
+  exhausted: boolean;
+  consecutive429: number;
+}
+
+/** Consecutive 429s within this window count toward exhaustion detection. */
+const CONSECUTIVE_WINDOW_MS = 2 * 60_000;
+/** When detected as exhausted, cooldown for this long. */
+const EXHAUSTED_COOLDOWN_MS = 2 * 3600_000;
+/** Retry-After threshold (seconds) above which we consider quota exhausted. */
+const EXHAUSTED_THRESHOLD_S = 300;
+/** Consecutive 429 count to trigger exhaustion detection. */
+const EXHAUSTED_CONSECUTIVE = 3;
+/** Default burst cooldown when no Retry-After header. */
+const DEFAULT_BURST_S = 30;
+
+const entries = new Map<string, CooldownEntry>();
+
+function key(pool: QuotaPool, account: number): string {
+  return `${pool}:${account}`;
+}
+
+export function isCoolingDown(pool: QuotaPool, account: number): boolean {
+  const entry = entries.get(key(pool, account));
+  if (!entry) return false;
+  if (Date.now() >= entry.until) {
+    entries.delete(key(pool, account));
+    return false;
+  }
+  return true;
+}
+
+export function isExhausted(pool: QuotaPool, account: number): boolean {
+  const entry = entries.get(key(pool, account));
+  if (!entry) return false;
+  if (Date.now() >= entry.until) {
+    entries.delete(key(pool, account));
+    return false;
+  }
+  return entry.exhausted;
+}
+
+export function record429(pool: QuotaPool, account: number, retryAfterSeconds?: number): void {
+  const k = key(pool, account);
+  const entry = entries.get(k) ?? { until: 0, exhausted: false, consecutive429: 0 };
+
+  entry.consecutive429++;
+  const retryAfter = retryAfterSeconds ?? DEFAULT_BURST_S;
+
+  if (retryAfter > EXHAUSTED_THRESHOLD_S || entry.consecutive429 >= EXHAUSTED_CONSECUTIVE) {
+    entry.exhausted = true;
+    entry.until = Date.now() + EXHAUSTED_COOLDOWN_MS;
+    logger.warn(`Quota exhausted: ${k}`, { cooldownMinutes: EXHAUSTED_COOLDOWN_MS / 60_000 });
+  } else {
+    entry.until = Date.now() + retryAfter * 1000;
+    logger.debug(`Burst cooldown: ${k}`, { retryAfterSeconds: retryAfter });
+  }
+
+  entries.set(k, entry);
+}
+
+export function recordSuccess(pool: QuotaPool, account: number): void {
+  const k = key(pool, account);
+  const entry = entries.get(k);
+  if (entry) {
+    entry.consecutive429 = 0;
+    entry.exhausted = false;
+    entries.delete(k);
+  }
+}
+
+/** Parse Retry-After header (seconds or HTTP-date). */
+export function parseRetryAfter(header: string | null): number | undefined {
+  if (!header) return undefined;
+  const seconds = Number(header);
+  if (!Number.isNaN(seconds)) return seconds;
+  const date = Date.parse(header);
+  if (!Number.isNaN(date)) return Math.max(0, Math.ceil((date - Date.now()) / 1000));
+  return undefined;
+}