ampcode-connector 0.1.0

package/src/routing/router.ts

@@ -0,0 +1,201 @@
+/**
+ * Route decision: local provider (free) or Amp upstream (paid).
+ *
+ * Multi-account aware:
+ * - 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)
+ */
+
+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 { logger, type RouteDecision } from "../utils/logger.ts";
+import * as affinity from "./affinity.ts";
+import type { QuotaPool } from "./cooldown.ts";
+import * as cooldown from "./cooldown.ts";
+
+export interface RouteResult {
+  decision: RouteDecision;
+  provider: string;
+  model: string;
+  handler: Provider | null;
+  account: number;
+  pool: QuotaPool | null;
+}
+
+interface Candidate {
+  provider: Provider;
+  pool: QuotaPool;
+  account: number;
+}
+
+export function routeRequest(
+  ampProvider: string,
+  model: string | null,
+  config: ProxyConfig,
+  threadId?: string,
+): RouteResult {
+  const modelStr = model ?? "unknown";
+
+  // Check thread affinity first
+  if (threadId) {
+    const pinned = affinity.get(threadId);
+    if (pinned && !cooldown.isExhausted(pinned.pool, pinned.account)) {
+      const handler = providerForPool(pinned.pool);
+      if (handler && handler.isAvailable(pinned.account)) {
+        if (!cooldown.isCoolingDown(pinned.pool, pinned.account)) {
+          logger.route(handler.routeDecision, ampProvider, modelStr);
+          return result(handler, ampProvider, modelStr, pinned.account, pinned.pool);
+        }
+        // Burst cooldown — still pinned but cooling, let it fall through to find alternative
+      }
+    }
+    // Affinity broken (exhausted / unavailable) — clear and re-route
+    if (pinned) affinity.clear(threadId);
+  }
+
+  // Build candidate list
+  const candidates = buildCandidates(ampProvider, config);
+  if (candidates.length === 0) {
+    logger.route("AMP_UPSTREAM", ampProvider, modelStr);
+    return result(null, ampProvider, modelStr, 0, null);
+  }
+
+  // Pick best candidate: not cooling down, least active threads
+  const picked = pickCandidate(candidates);
+  if (!picked) {
+    logger.route("AMP_UPSTREAM", ampProvider, modelStr);
+    return result(null, ampProvider, modelStr, 0, null);
+  }
+
+  // Pin thread affinity
+  if (threadId) affinity.set(threadId, picked.pool, picked.account);
+
+  logger.route(picked.provider.routeDecision, ampProvider, modelStr);
+  return result(picked.provider, ampProvider, modelStr, picked.account, picked.pool);
+}
+
+/** Record a 429 response and attempt re-route. Returns a new RouteResult or null. */
+export function rerouteAfter429(
+  ampProvider: string,
+  model: string | null,
+  config: ProxyConfig,
+  failedPool: QuotaPool,
+  failedAccount: number,
+  retryAfterSeconds: number | undefined,
+  threadId?: string,
+): RouteResult | null {
+  cooldown.record429(failedPool, failedAccount, retryAfterSeconds);
+
+  // If exhausted, break thread affinity
+  if (threadId && cooldown.isExhausted(failedPool, failedAccount)) {
+    affinity.clear(threadId);
+  }
+
+  const modelStr = model ?? "unknown";
+  const candidates = buildCandidates(ampProvider, config);
+  const picked = pickCandidate(candidates);
+
+  if (!picked) return null;
+
+  if (threadId) affinity.set(threadId, picked.pool, picked.account);
+  logger.route(picked.provider.routeDecision, ampProvider, modelStr);
+  return result(picked.provider, ampProvider, modelStr, picked.account, picked.pool);
+}
+
+/** Record a successful response — clears cooldown. */
+export function recordSuccess(pool: QuotaPool, account: number): void {
+  cooldown.recordSuccess(pool, account);
+}
+
+function buildCandidates(ampProvider: string, config: ProxyConfig): Candidate[] {
+  const candidates: Candidate[] = [];
+
+  switch (ampProvider) {
+    case "anthropic":
+      if (config.providers.anthropic) {
+        addAccountCandidates(candidates, anthropic, "anthropic", "anthropic");
+      }
+      break;
+
+    case "openai":
+      if (config.providers.codex) {
+        addAccountCandidates(candidates, codex, "codex", "codex");
+      }
+      break;
+
+    case "google":
+      if (config.providers.google) {
+        // Both gemini and antigravity use "google" credentials — separate quota pools
+        addAccountCandidates(candidates, gemini, "gemini", "google");
+        addAccountCandidates(candidates, antigravity, "antigravity", "google");
+      }
+      break;
+  }
+
+  return candidates;
+}
+
+function addAccountCandidates(
+  candidates: Candidate[],
+  provider: Provider,
+  pool: QuotaPool,
+  providerName: ProviderName,
+): void {
+  for (const { account, credentials } of store.getAll(providerName)) {
+    if (credentials.refreshToken) {
+      candidates.push({ provider, pool, account });
+    }
+  }
+}
+
+function pickCandidate(candidates: Candidate[]): Candidate | null {
+  // Filter out cooling-down accounts
+  const available = candidates.filter((c) => !cooldown.isCoolingDown(c.pool, c.account));
+  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);
+
+  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;
+    }
+  }
+
+  return best;
+}
+
+function providerForPool(pool: QuotaPool): Provider | null {
+  switch (pool) {
+    case "anthropic":
+      return anthropic;
+    case "codex":
+      return codex;
+    case "gemini":
+      return gemini;
+    case "antigravity":
+      return antigravity;
+  }
+}
+
+function result(
+  handler: Provider | null,
+  provider: string,
+  model: string,
+  account: number,
+  pool: QuotaPool | null,
+): RouteResult {
+  const decision: RouteDecision = handler?.routeDecision ?? "AMP_UPSTREAM";
+  return { decision, provider, model, handler, account, pool };
+}

package/src/server/server.ts

@@ -0,0 +1,147 @@
+/** HTTP server — routes provider requests through local OAuth or Amp upstream. */
+
+import type { ProxyConfig } from "../config/config.ts";
+import * as rewriter from "../proxy/rewriter.ts";
+import * as upstream from "../proxy/upstream.ts";
+import { parseRetryAfter, record429 } from "../routing/cooldown.ts";
+import { recordSuccess, rerouteAfter429, routeRequest } from "../routing/router.ts";
+import { logger } from "../utils/logger.ts";
+import * as path from "../utils/path.ts";
+
+/** Max 429-reroute attempts before falling back to upstream. */
+const MAX_REROUTE_ATTEMPTS = 4;
+
+export function startServer(config: ProxyConfig): ReturnType<typeof Bun.serve> {
+  const server = Bun.serve({
+    port: config.port,
+    hostname: "localhost",
+
+    async fetch(req) {
+      const startTime = Date.now();
+      try {
+        const response = await handle(req, config);
+        logRequest(req, startTime, response.status);
+        return response;
+      } catch (err) {
+        logRequest(req, startTime, 500);
+        logger.error("Unhandled server error", { error: String(err) });
+        return new Response(JSON.stringify({ error: "Internal proxy error" }), {
+          status: 500,
+          headers: { "Content-Type": "application/json" },
+        });
+      }
+    },
+  });
+
+  logger.info(`ampcode-connector listening on http://localhost:${config.port}`);
+  return server;
+}
+
+async function handle(req: Request, config: ProxyConfig): Promise<Response> {
+  const { pathname } = new URL(req.url);
+
+  if ((pathname === "/" || pathname === "/status") && req.method === "GET") {
+    return healthCheck(config);
+  }
+
+  if (path.browser(pathname)) {
+    const target = new URL(pathname + new URL(req.url).search, config.ampUpstreamUrl);
+    return Response.redirect(target.toString(), 302);
+  }
+
+  if (path.passthrough(pathname)) {
+    return upstream.forward(req, config.ampUpstreamUrl, config.ampApiKey);
+  }
+
+  const providerName = path.provider(pathname);
+  if (providerName) return handleProvider(req, providerName, pathname, config);
+
+  return upstream.forward(req, config.ampUpstreamUrl, config.ampApiKey);
+}
+
+async function handleProvider(
+  req: Request,
+  providerName: string,
+  pathname: string,
+  config: ProxyConfig,
+): Promise<Response> {
+  const sub = path.subpath(pathname);
+  const threadId = req.headers.get("x-amp-thread-id") ?? undefined;
+
+  let body = "";
+  if (req.method === "POST") body = await req.text();
+
+  const model = path.model(body) ?? path.modelFromUrl(sub);
+  let route = routeRequest(providerName, model, config, threadId);
+
+  logger.info(`ROUTE ${route.decision} provider=${providerName} model=${model ?? "?"} account=${route.account} sub=${sub}`);
+
+  if (route.handler) {
+    const rewrite = model ? rewriter.rewrite(model) : undefined;
+    const response = await route.handler.forward(sub, body, req.headers, rewrite, route.account);
+
+    // 429 → attempt reroute to different account/pool
+    if (response.status === 429 && route.pool) {
+      const retryAfter = parseRetryAfter(response.headers.get("retry-after"));
+      logger.warn(`429 from ${route.decision} account=${route.account}`, { retryAfter });
+
+      for (let attempt = 0; attempt < MAX_REROUTE_ATTEMPTS; attempt++) {
+        const next = rerouteAfter429(providerName, model, config, route.pool, route.account, retryAfter, threadId);
+        if (!next) break;
+
+        route = next;
+        logger.info(`REROUTE → ${route.decision} account=${route.account}`);
+        const retryResponse = await route.handler!.forward(sub, body, req.headers, rewrite, route.account);
+
+        if (retryResponse.status === 429 && route.pool) {
+          const nextRetryAfter = parseRetryAfter(retryResponse.headers.get("retry-after"));
+          record429(route.pool, route.account, nextRetryAfter);
+          continue;
+        }
+
+        if (retryResponse.status !== 401) {
+          if (route.pool) recordSuccess(route.pool, route.account);
+          return retryResponse;
+        }
+        break;
+      }
+
+      // All reroutes failed — fall through to upstream
+    } else if (response.status === 401) {
+      logger.debug("Local provider denied, falling back to upstream");
+    } else {
+      if (route.pool) recordSuccess(route.pool, route.account);
+      return response;
+    }
+  }
+
+  const upstreamReq = new Request(req.url, {
+    method: req.method,
+    headers: req.headers,
+    body: body || undefined,
+  });
+  return upstream.forward(upstreamReq, config.ampUpstreamUrl, config.ampApiKey);
+}
+
+function healthCheck(config: ProxyConfig): Response {
+  return new Response(
+    JSON.stringify(
+      {
+        status: "ok",
+        service: "ampcode-connector",
+        port: config.port,
+        upstream: config.ampUpstreamUrl,
+        providers: config.providers,
+      },
+      null,
+      2,
+    ),
+    { status: 200, headers: { "Content-Type": "application/json" } },
+  );
+}
+
+function logRequest(req: Request, startTime: number, statusCode: number): void {
+  const duration = Date.now() - startTime;
+  const { pathname } = new URL(req.url);
+  logger.info(`${req.method} ${pathname}`, { duration });
+}

package/src/utils/browser.ts

@@ -0,0 +1,19 @@
+/** Cross-platform: macOS (open), Linux (xdg-open), Windows (start). */
+export async function open(url: string): Promise<boolean> {
+  const commands: Record<string, string[]> = {
+    darwin: ["open", url],
+    linux: ["xdg-open", url],
+    win32: ["cmd", "/c", "start", url],
+  };
+
+  const cmd = commands[process.platform];
+  if (!cmd) return false;
+
+  try {
+    const proc = Bun.spawn(cmd);
+    await proc.exited;
+    return true;
+  } catch {
+    return false;
+  }
+}

package/src/utils/code-assist.ts

@@ -0,0 +1,42 @@
+/** Cloud Code Assist request/URL helpers shared by Gemini CLI and Antigravity. */
+
+interface WrapOptions {
+  projectId: string;
+  model: string;
+  body: Record<string, unknown>;
+  userAgent: "antigravity" | "pi-coding-agent";
+  requestIdPrefix: "agent" | "pi";
+  requestType?: "agent";
+}
+
+/** Wrap a raw request body in the Cloud Code Assist envelope. */
+export function wrapRequest(opts: WrapOptions): string {
+  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)}`,
+  });
+}
+
+/** Build the Cloud Code Assist URL for a given action. Preserves the original action
+ *  (generateContent vs streamGenerateContent). Only adds ?alt=sse for streaming actions. */
+export function buildUrl(endpoint: string, action: string): string {
+  const streaming = action.toLowerCase().includes("stream");
+  return `${endpoint}/v1internal:${action}${streaming ? "?alt=sse" : ""}`;
+}
+
+/** Unwrap Cloud Code Assist SSE envelope: {"response":{...},"traceId":"..."} → inner response.
+ *  Returns empty string for [DONE] sentinel (Google SDK doesn't expect it). */
+export function unwrap(data: string): string {
+  if (data === "[DONE]") return "";
+  try {
+    const parsed = JSON.parse(data) as { response?: unknown };
+    if (parsed.response !== undefined) return JSON.stringify(parsed.response);
+    return data;
+  } catch {
+    return data;
+  }
+}

package/src/utils/encoding.ts

@@ -0,0 +1,9 @@
+/** Encode bytes to base64url (no padding). */
+export function toBase64url(buffer: Uint8Array): string {
+  return Buffer.from(buffer).toString("base64url");
+}
+
+/** Decode base64url string to bytes. Handles both base64url and standard base64. */
+export function fromBase64url(input: string): Uint8Array {
+  return new Uint8Array(Buffer.from(input, "base64url"));
+}

package/src/utils/logger.ts

@@ -0,0 +1,80 @@
+/** Structured logging with route decision tracking. */
+
+export type RouteDecision = "LOCAL_CLAUDE" | "LOCAL_CODEX" | "LOCAL_GEMINI" | "LOCAL_ANTIGRAVITY" | "AMP_UPSTREAM";
+
+export type LogLevel = "debug" | "info" | "warn" | "error";
+
+interface LogEntry {
+  timestamp: string;
+  level: LogLevel;
+  message: string;
+  route?: RouteDecision;
+  provider?: string;
+  model?: string;
+  duration?: number;
+  error?: string;
+}
+
+const LOG_LEVELS: Record<LogLevel, number> = {
+  debug: 0,
+  info: 1,
+  warn: 2,
+  error: 3,
+};
+
+let currentLevel: LogLevel = "info";
+
+export function setLogLevel(level: LogLevel): void {
+  currentLevel = level;
+}
+
+function shouldLog(level: LogLevel): boolean {
+  return LOG_LEVELS[level] >= LOG_LEVELS[currentLevel];
+}
+
+function format(entry: LogEntry): string {
+  const { timestamp, level, message, route, provider, model, duration, error } = entry;
+
+  let line = `${timestamp} [${level.toUpperCase().padEnd(5)}] ${message}`;
+  if (route) line += ` route=${route}`;
+  if (provider) line += ` provider=${provider}`;
+  if (model) line += ` model=${model}`;
+  if (duration !== undefined) line += ` duration=${duration}ms`;
+  if (error) line += ` error=${error}`;
+
+  return line;
+}
+
+type Meta = Partial<Omit<LogEntry, "timestamp" | "level" | "message">>;
+
+function log(level: LogLevel, message: string, meta?: Meta): void {
+  if (!shouldLog(level)) return;
+
+  const entry: LogEntry = {
+    timestamp: new Date().toISOString(),
+    level,
+    message,
+    ...meta,
+  };
+
+  const line = format(entry);
+
+  if (level === "error") {
+    console.error(line);
+  } else if (level === "warn") {
+    console.warn(line);
+  } else {
+    console.log(line);
+  }
+}
+
+export const logger = {
+  debug: (message: string, meta?: Meta) => log("debug", message, meta),
+  info: (message: string, meta?: Meta) => log("info", message, meta),
+  warn: (message: string, meta?: Meta) => log("warn", message, meta),
+  error: (message: string, meta?: Meta) => log("error", message, meta),
+
+  route(decision: RouteDecision, provider: string, model: string): void {
+    log("info", "Route decision", { route: decision, provider, model });
+  },
+};

package/src/utils/path.ts

@@ -0,0 +1,52 @@
+/** URL path parsing for Amp CLI provider routes. */
+
+import { browserPrefixes, passthroughExact, passthroughPrefixes } from "../constants.ts";
+
+export function passthrough(pathname: string): boolean {
+  if ((passthroughExact as readonly string[]).includes(pathname)) return true;
+  return passthroughPrefixes.some((prefix) => pathname.startsWith(prefix));
+}
+
+export function browser(pathname: string): boolean {
+  if ((passthroughExact as readonly string[]).includes(pathname)) return true;
+  return browserPrefixes.some((prefix) => pathname.startsWith(prefix));
+}
+
+export function provider(pathname: string): string | null {
+  const match = pathname.match(/^\/api\/provider\/([^/]+)/);
+  return match?.[1] ?? null;
+}
+
+export function subpath(pathname: string): string {
+  const match = pathname.match(/^\/api\/provider\/[^/]+(\/.*)/);
+  return match?.[1] ?? pathname;
+}
+
+export function model(body: string): string | null {
+  try {
+    const parsed = JSON.parse(body) as { model?: string };
+    return parsed.model ?? null;
+  } catch {
+    return null;
+  }
+}
+
+export function modelFromUrl(url: string): string | null {
+  const match = url.match(/models\/([^/:]+)/);
+  return match?.[1] ?? null;
+}
+
+export function gemini(url: string): { model: string; action: string } | null {
+  const match = url.match(/models\/([^/:]+):(\w+)/);
+  if (!match) return null;
+  return { model: match[1]!, action: match[2]! };
+}
+
+export function streaming(body: string): boolean {
+  try {
+    const parsed = JSON.parse(body) as { stream?: boolean };
+    return parsed.stream === true;
+  } catch {
+    return false;
+  }
+}

package/src/utils/streaming.ts

@@ -0,0 +1,124 @@
+/** SSE (Server-Sent Events) stream parsing, encoding, and transformation. */
+
+export interface Chunk {
+  event?: string;
+  data: string;
+  id?: string;
+  retry?: number;
+}
+
+export function parse(raw: string): Chunk[] {
+  const chunks: Chunk[] = [];
+
+  for (const block of raw.split("\n\n")) {
+    if (!block.trim()) continue;
+
+    const chunk: Partial<Chunk> = {};
+    for (const line of block.split("\n")) {
+      if (line.startsWith("event:")) {
+        chunk.event = line.slice(6).trim();
+      } else if (line.startsWith("data:")) {
+        const value = line.slice(5).trimStart();
+        chunk.data = chunk.data ? `${chunk.data}\n${value}` : value;
+      } else if (line.startsWith("id:")) {
+        chunk.id = line.slice(3).trim();
+      } else if (line.startsWith("retry:")) {
+        const val = parseInt(line.slice(6).trim(), 10);
+        if (!isNaN(val)) chunk.retry = val;
+      }
+    }
+
+    if (chunk.data !== undefined) chunks.push(chunk as Chunk);
+  }
+
+  return chunks;
+}
+
+export function encode(chunk: Chunk): string {
+  let result = "";
+  if (chunk.event) result += `event: ${chunk.event}\n`;
+  if (chunk.id) result += `id: ${chunk.id}\n`;
+  if (chunk.retry !== undefined) result += `retry: ${chunk.retry}\n`;
+
+  for (const line of chunk.data.split("\n")) {
+    result += `data: ${line}\n`;
+  }
+
+  result += "\n";
+  return result;
+}
+
+export function transform(
+  source: ReadableStream<Uint8Array>,
+  fn: (data: string) => string,
+): ReadableStream<Uint8Array> {
+  const decoder = new TextDecoder();
+  const encoder = new TextEncoder();
+  let buffer = "";
+
+  const stream = new TransformStream<Uint8Array, Uint8Array>({
+    transform(raw, controller) {
+      buffer += decoder.decode(raw, { stream: true }).replaceAll("\r\n", "\n");
+
+      const boundary = buffer.lastIndexOf("\n\n");
+      if (boundary === -1) return;
+
+      const complete = buffer.slice(0, boundary + 2);
+      buffer = buffer.slice(boundary + 2);
+
+      for (const chunk of parse(complete)) {
+        chunk.data = fn(chunk.data);
+        if (chunk.data) controller.enqueue(encoder.encode(encode(chunk)));
+      }
+    },
+
+    flush(controller) {
+      if (buffer.trim()) {
+        for (const chunk of parse(buffer)) {
+          chunk.data = fn(chunk.data);
+          if (chunk.data) controller.enqueue(encoder.encode(encode(chunk)));
+        }
+      }
+      buffer = "";
+    },
+  });
+
+  return source.pipeThrough(stream);
+}
+
+export function headers(): Record<string, string> {
+  return {
+    "Content-Type": "text/event-stream",
+    "Cache-Control": "no-cache",
+    Connection: "keep-alive",
+  };
+}
+
+const forwardedHeaders = [
+  "x-request-id",
+  "request-id",
+  "anthropic-ratelimit-requests-limit",
+  "anthropic-ratelimit-requests-remaining",
+  "anthropic-ratelimit-tokens-limit",
+  "anthropic-ratelimit-tokens-remaining",
+  "x-ratelimit-limit-requests",
+  "x-ratelimit-remaining-requests",
+  "x-ratelimit-limit-tokens",
+  "x-ratelimit-remaining-tokens",
+] as const;
+
+export function proxy(upstream: Response, rewrite?: (data: string) => string): Response {
+  if (!upstream.body) {
+    return new Response("No response body", { status: 502 });
+  }
+
+  const body = rewrite ? transform(upstream.body, rewrite) : upstream.body;
+
+  const h: Record<string, string> = { ...headers() };
+  for (const name of forwardedHeaders) {
+    const value = upstream.headers.get(name);
+    if (value) h[name] = value;
+  }
+
+  return new Response(body, { status: upstream.status, headers: h });
+}