ampcode-connector 0.1.6 → 0.1.7

package/package.json

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

package/src/auth/store.ts

@@ -18,8 +18,8 @@ export interface Credentials {
 
 export type ProviderName = "anthropic" | "codex" | "google";
 
-const DIR = join(homedir(), ".ampcode-connector");
-const DB_PATH = join(DIR, "credentials.db");
+const DEFAULT_DIR = join(homedir(), ".ampcode-connector");
+const DEFAULT_DB_PATH = join(DEFAULT_DIR, "credentials.db");
 
 interface DataRow {
   data: string;
@@ -47,12 +47,19 @@ interface Statements {
 
 let _db: Database | null = null;
 let _stmts: Statements | null = null;
+let _dbPath = DEFAULT_DB_PATH;
+
+/** Override the database path (must be called before any store operation). */
+export function setDbPath(path: string): void {
+  _dbPath = path;
+}
 
 function init() {
   if (_stmts) return _stmts;
 
-  mkdirSync(DIR, { recursive: true, mode: 0o700 });
-  _db = new Database(DB_PATH, { strict: true });
+  const dir = _dbPath.replace(/\/[^/]+$/, "");
+  mkdirSync(dir, { recursive: true, mode: 0o700 });
+  _db = new Database(_dbPath, { strict: true });
   _db.exec("PRAGMA journal_mode=WAL");
   _db.exec("PRAGMA busy_timeout=5000");
   _db.exec(`

package/src/cli/ads.ts

@@ -0,0 +1,38 @@
+/** Periodic GitHub star reminder — non-intrusive, shows in server logs. */
+
+import { line, s } from "../cli/ansi.ts";
+
+const REPO_URL = "https://github.com/nghyane/ampcode-connector";
+const REQUEST_INTERVAL = 50;
+
+let requestCount = 0;
+let shown = false;
+
+const messages = [
+  `${s.yellow}⭐${s.reset} Enjoying ampcode-connector? Star us on GitHub → ${s.cyan}${REPO_URL}${s.reset}`,
+  `${s.yellow}⭐${s.reset} Help others discover this tool — star on GitHub → ${s.cyan}${REPO_URL}${s.reset}`,
+  `${s.yellow}⭐${s.reset} ${s.dim}Your star helps keep this project alive!${s.reset} → ${s.cyan}${REPO_URL}${s.reset}`,
+];
+
+function pick(): string {
+  return messages[Math.floor(Math.random() * messages.length)]!;
+}
+
+/** Show star prompt in the startup banner (once). */
+export function bannerAd(): void {
+  line(`  ${s.dim}⭐ Star us → ${REPO_URL}${s.reset}`);
+}
+
+/** Call after each proxied request. Shows a reminder every N requests. */
+export function maybeShowAd(): void {
+  requestCount++;
+  if (requestCount % REQUEST_INTERVAL !== 0) return;
+
+  // Only show once per interval, don't spam
+  if (shown && requestCount < REQUEST_INTERVAL * 3) return;
+  shown = true;
+
+  line();
+  line(`  ${pick()}`);
+  line();
+}

package/src/constants.ts

@@ -25,9 +25,11 @@ export const codexHeaderValues = {
   USER_AGENT: `codex_cli_rs/0.101.0 (${process.platform} ${process.arch})`,
 } as const;
 
-/** Map /v1/responses → /codex/responses for the ChatGPT backend. */
+/** Map Amp CLI paths → ChatGPT backend paths.
+ *  Both /v1/responses and /v1/chat/completions route to /codex/responses. */
 export const codexPathMap: Record<string, string> = {
   "/v1/responses": "/codex/responses",
+  "/v1/chat/completions": "/codex/responses",
 } as const;
 export const DEFAULT_AMP_UPSTREAM_URL = "https://ampcode.com";
 

package/src/index.ts

@@ -5,6 +5,7 @@ import { startAutoRefresh } from "./auth/auto-refresh.ts";
 import * as configs from "./auth/configs.ts";
 import type { OAuthConfig } from "./auth/oauth.ts";
 import * as oauth from "./auth/oauth.ts";
+import { bannerAd } from "./cli/ads.ts";
 import { line, s } from "./cli/ansi.ts";
 import { setup } from "./cli/setup.ts";
 import * as status from "./cli/status.ts";
@@ -72,6 +73,8 @@ function banner(config: ProxyConfig): void {
   line();
   line(`  ${s.dim}upstream → ${upstream}${s.reset}`);
   line();
+  bannerAd();
+  line();
 }
 
 function usage(): void {

package/src/providers/codex-sse.ts

@@ -0,0 +1,247 @@
+/** Transforms Responses API SSE events → Chat Completions SSE chunks.
+ *
+ *  Codex backend returns Responses API format (response.output_text.delta, etc.)
+ *  but Amp CLI expects Chat Completions format (chat.completion.chunk). */
+
+import * as sse from "../utils/streaming.ts";
+
+interface CompletionChunk {
+  id: string;
+  object: "chat.completion.chunk";
+  created: number;
+  model: string;
+  choices: Choice[];
+  usage?: Usage | null;
+}
+
+interface Choice {
+  index: number;
+  delta: Delta;
+  finish_reason: string | null;
+}
+
+interface Delta {
+  role?: string;
+  content?: string;
+  tool_calls?: ToolCallDelta[];
+}
+
+interface ToolCallDelta {
+  index: number;
+  id?: string;
+  type?: string;
+  function?: { name?: string; arguments?: string };
+}
+
+interface Usage {
+  prompt_tokens: number;
+  completion_tokens: number;
+  total_tokens: number;
+  prompt_tokens_details?: { cached_tokens: number };
+}
+
+interface TransformState {
+  responseId: string;
+  model: string;
+  created: number;
+  toolCallIndex: number;
+  /** Track active tool call IDs to assign sequential indices. */
+  toolCallIds: Map<string, number>;
+}
+
+/** Create a stateful SSE transformer: Responses API → Chat Completions. */
+export function createResponseTransformer(ampModel: string): (data: string) => string {
+  const state: TransformState = {
+    responseId: "",
+    model: ampModel,
+    created: Math.floor(Date.now() / 1000),
+    toolCallIndex: 0,
+    toolCallIds: new Map(),
+  };
+
+  return (data: string): string => {
+    if (data === "[DONE]") return data;
+
+    let parsed: Record<string, unknown>;
+    try {
+      parsed = JSON.parse(data) as Record<string, unknown>;
+    } catch {
+      return data;
+    }
+
+    const eventType = parsed.type as string | undefined;
+    if (!eventType) return data;
+
+    // Extract response metadata on creation
+    if (eventType === "response.created") {
+      const resp = parsed.response as Record<string, unknown>;
+      state.responseId = (resp?.id as string) ?? state.responseId;
+      state.model = ampModel;
+      state.created = (resp?.created_at as number) ?? state.created;
+      // Don't emit a chunk for response.created
+      return "";
+    }
+
+    switch (eventType) {
+      // Assistant message started — emit role
+      case "response.output_item.added": {
+        const item = parsed.item as Record<string, unknown>;
+        if (item?.type === "message" && item.role === "assistant") {
+          return serialize(state, { role: "assistant", content: "" });
+        }
+        if (item?.type === "function_call") {
+          const callId = item.call_id as string;
+          const name = item.name as string;
+          const idx = state.toolCallIndex++;
+          state.toolCallIds.set(callId, idx);
+          return serialize(state, {
+            tool_calls: [{ index: idx, id: callId, type: "function", function: { name, arguments: "" } }],
+          });
+        }
+        return "";
+      }
+
+      // Text content delta
+      case "response.output_text.delta": {
+        const delta = parsed.delta as string;
+        if (delta) return serialize(state, { content: delta });
+        return "";
+      }
+
+      // Function call arguments delta
+      case "response.function_call_arguments.delta": {
+        const delta = parsed.delta as string;
+        const callId = parsed.call_id as string | undefined;
+        if (delta) {
+          const idx = callId ? (state.toolCallIds.get(callId) ?? 0) : 0;
+          return serialize(state, { tool_calls: [{ index: idx, function: { arguments: delta } }] });
+        }
+        return "";
+      }
+
+      // Response completed — emit finish_reason + usage
+      case "response.completed": {
+        const resp = parsed.response as Record<string, unknown>;
+        const usage = extractUsage(resp?.usage as Record<string, unknown> | undefined);
+        const hasToolCalls = state.toolCallIndex > 0;
+        const finishReason = hasToolCalls ? "tool_calls" : "stop";
+        return serializeFinish(state, finishReason, usage);
+      }
+
+      // Reasoning/thinking delta — emit as content (Amp shows thinking)
+      case "response.reasoning_summary_text.delta": {
+        const delta = parsed.delta as string;
+        if (delta) return serialize(state, { content: delta });
+        return "";
+      }
+
+      // Events we can skip
+      case "response.in_progress":
+      case "response.output_item.done":
+      case "response.content_part.added":
+      case "response.content_part.done":
+      case "response.output_text.done":
+      case "response.function_call_arguments.done":
+      case "response.reasoning_summary_part.added":
+      case "response.reasoning_summary_part.done":
+        return "";
+
+      default:
+        return "";
+    }
+  };
+}
+
+function serialize(state: TransformState, delta: Delta): string {
+  const chunk: CompletionChunk = {
+    id: `chatcmpl-${state.responseId}`,
+    object: "chat.completion.chunk",
+    created: state.created,
+    model: state.model,
+    choices: [{ index: 0, delta, finish_reason: null }],
+  };
+  return JSON.stringify(chunk);
+}
+
+function serializeFinish(state: TransformState, finishReason: string, usage?: Usage): string {
+  const chunk: CompletionChunk = {
+    id: `chatcmpl-${state.responseId}`,
+    object: "chat.completion.chunk",
+    created: state.created,
+    model: state.model,
+    choices: [{ index: 0, delta: {}, finish_reason: finishReason }],
+    ...(usage ? { usage } : {}),
+  };
+  return JSON.stringify(chunk);
+}
+
+function extractUsage(raw: Record<string, unknown> | undefined): Usage | undefined {
+  if (!raw) return undefined;
+  const input = (raw.input_tokens as number) ?? 0;
+  const output = (raw.output_tokens as number) ?? 0;
+  const cached = (raw.input_tokens_details as Record<string, unknown>)?.cached_tokens as number | undefined;
+  return {
+    prompt_tokens: input,
+    completion_tokens: output,
+    total_tokens: input + output,
+    ...(cached !== undefined ? { prompt_tokens_details: { cached_tokens: cached } } : {}),
+  };
+}
+
+/** Wrap a Codex SSE response with the Responses → Chat Completions transformer.
+ *  Strips Responses API event names so output looks like standard Chat Completions SSE. */
+export function transformCodexResponse(response: Response, ampModel: string): Response {
+  if (!response.body) return response;
+
+  const transformer = createResponseTransformer(ampModel);
+  const body = transformStream(response.body, transformer);
+
+  return new Response(body, {
+    status: response.status,
+    headers: {
+      "Content-Type": "text/event-stream",
+      "Cache-Control": "no-cache",
+      Connection: "keep-alive",
+    },
+  });
+}
+
+/** Custom SSE transform that strips event names (Chat Completions doesn't use them). */
+function transformStream(source: ReadableStream<Uint8Array>, fn: (data: string) => string): ReadableStream<Uint8Array> {
+  const decoder = new TextDecoder();
+  const textEncoder = new TextEncoder();
+  let buffer = "";
+
+  return source.pipeThrough(
+    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 sse.parse(complete)) {
+          const transformed = fn(chunk.data);
+          if (transformed) {
+            // Emit without event name — standard Chat Completions format
+            controller.enqueue(textEncoder.encode(`data: ${transformed}\n\n`));
+          }
+        }
+      },
+      flush(controller) {
+        if (buffer.trim()) {
+          for (const chunk of sse.parse(buffer)) {
+            const transformed = fn(chunk.data);
+            if (transformed) {
+              controller.enqueue(textEncoder.encode(`data: ${transformed}\n\n`));
+            }
+          }
+        }
+        // Emit [DONE] marker
+        controller.enqueue(textEncoder.encode("data: [DONE]\n\n"));
+      },
+    }),
+  );
+}

package/src/providers/codex.ts

@@ -1,7 +1,8 @@
 /** Forwards requests to chatgpt.com/backend-api/codex with Codex CLI OAuth token.
  *
- *  The ChatGPT backend requires specific headers (account-id from JWT, originator,
- *  OpenAI-Beta) and a different URL path (/codex/responses) than api.openai.com. */
+ *  The ChatGPT backend only accepts the Responses API format (input[] + instructions),
+ *  but Amp CLI sends Chat Completions format (messages[]). This module transforms
+ *  the request body before forwarding. */
 
 import { codex as config } from "../auth/configs.ts";
 import * as oauth from "../auth/oauth.ts";
@@ -10,6 +11,9 @@ import { CODEX_BASE_URL, codexHeaders, codexHeaderValues, codexPathMap } from ".
 import { fromBase64url } from "../utils/encoding.ts";
 import type { Provider } from "./base.ts";
 import { denied, forward } from "./base.ts";
+import { transformCodexResponse } from "./codex-sse.ts";
+
+const DEFAULT_INSTRUCTIONS = "You are an expert coding assistant.";
 
 export const provider: Provider = {
   name: "OpenAI Codex",
@@ -26,13 +30,16 @@ export const provider: Provider = {
 
     const accountId = getAccountId(accessToken, account);
     const codexPath = codexPathMap[sub] ?? sub;
+    const { body: codexBody, needsResponseTransform } = transformForCodex(body.forwardBody);
+    const ampModel = body.ampModel ?? "gpt-5.2";
 
-    return forward({
+    const response = await forward({
       url: `${CODEX_BASE_URL}${codexPath}`,
-      body: body.forwardBody,
+      body: codexBody,
       streaming: body.stream,
       providerName: "OpenAI Codex",
-      rewrite,
+      // Skip generic rewrite when we need full response transform
+      rewrite: needsResponseTransform ? undefined : rewrite,
       headers: {
         "Content-Type": "application/json",
         Authorization: `Bearer ${accessToken}`,
@@ -45,9 +52,205 @@ export const provider: Provider = {
         ...(accountId ? { [codexHeaders.ACCOUNT_ID]: accountId } : {}),
       },
     });
+
+    // Transform Responses API SSE → Chat Completions SSE when original was messages[] format
+    if (needsResponseTransform && response.ok) {
+      return transformCodexResponse(response, ampModel);
+    }
+    return response;
   },
 };
 
+// ---------------------------------------------------------------------------
+// Body transformation: Chat Completions → Responses API
+// ---------------------------------------------------------------------------
+
+interface ChatMessage {
+  role: string;
+  content: unknown;
+  tool_calls?: ToolCallItem[];
+  tool_call_id?: string;
+  name?: string;
+}
+
+interface ToolCallItem {
+  id: string;
+  type: string;
+  function: { name: string; arguments: string };
+}
+
+function transformForCodex(rawBody: string): { body: string; needsResponseTransform: boolean } {
+  if (!rawBody) return { body: rawBody, needsResponseTransform: false };
+
+  let parsed: Record<string, unknown>;
+  try {
+    parsed = JSON.parse(rawBody) as Record<string, unknown>;
+  } catch {
+    return { body: rawBody, needsResponseTransform: false };
+  }
+
+  // Convert Chat Completions messages[] → Responses API input[]
+  let needsResponseTransform = false;
+  if (Array.isArray(parsed.messages) && !parsed.input) {
+    const { instructions, input } = convertMessages(parsed.messages as ChatMessage[]);
+    parsed.input = input;
+    parsed.instructions = parsed.instructions ?? instructions ?? DEFAULT_INSTRUCTIONS;
+    delete parsed.messages;
+    needsResponseTransform = true;
+  }
+
+  // Already Responses API format — ensure instructions exists
+  if (!parsed.instructions) {
+    parsed.instructions = extractInstructionsFromInput(parsed) ?? DEFAULT_INSTRUCTIONS;
+  }
+
+  // Codex backend requirements
+  parsed.store = false;
+  parsed.stream = true;
+
+  // Strip id fields from input items
+  if (Array.isArray(parsed.input)) {
+    stripInputIds(parsed.input as Record<string, unknown>[]);
+  }
+
+  // Remove fields the Codex backend doesn't accept
+  delete parsed.max_tokens;
+  delete parsed.max_completion_tokens;
+  // Chat Completions fields not in Responses API
+  delete parsed.frequency_penalty;
+  delete parsed.logprobs;
+  delete parsed.top_logprobs;
+  delete parsed.n;
+  delete parsed.presence_penalty;
+  delete parsed.seed;
+  delete parsed.stop;
+  delete parsed.logit_bias;
+  delete parsed.response_format;
+
+  return { body: JSON.stringify(parsed), needsResponseTransform };
+}
+
+/** Convert Chat Completions messages[] → Responses API input[] + instructions. */
+function convertMessages(messages: ChatMessage[]): { instructions: string | null; input: unknown[] } {
+  let instructions: string | null = null;
+  const input: unknown[] = [];
+
+  for (const msg of messages) {
+    switch (msg.role) {
+      case "system":
+      case "developer": {
+        // First system message → instructions; additional ones → developer input items
+        const text = textOf(msg.content);
+        if (!instructions) {
+          instructions = text;
+        } else if (text) {
+          input.push({ role: "developer", content: [{ type: "input_text", text }] });
+        }
+        break;
+      }
+
+      case "user":
+        input.push({ role: "user", content: convertUserContent(msg.content) });
+        break;
+
+      case "assistant": {
+        // Text content → message output item
+        const text = textOf(msg.content);
+        if (text) {
+          input.push({
+            type: "message",
+            role: "assistant",
+            content: [{ type: "output_text", text, annotations: [] }],
+            status: "completed",
+          });
+        }
+        // Tool calls → function_call items
+        if (msg.tool_calls) {
+          for (const tc of msg.tool_calls) {
+            input.push({
+              type: "function_call",
+              call_id: tc.id,
+              name: tc.function.name,
+              arguments: tc.function.arguments,
+            });
+          }
+        }
+        break;
+      }
+
+      case "tool":
+        // Tool result → function_call_output
+        input.push({
+          type: "function_call_output",
+          call_id: msg.tool_call_id,
+          output: textOf(msg.content) ?? "",
+        });
+        break;
+    }
+  }
+
+  return { instructions, input };
+}
+
+/** Convert user message content to Responses API format. */
+function convertUserContent(content: unknown): unknown[] {
+  if (typeof content === "string") {
+    return [{ type: "input_text", text: content }];
+  }
+  if (Array.isArray(content)) {
+    return content.map((part: Record<string, unknown>) => {
+      if (part.type === "text") {
+        return { type: "input_text", text: part.text };
+      }
+      if (part.type === "image_url") {
+        const imageUrl = part.image_url as Record<string, unknown>;
+        return { type: "input_image", image_url: imageUrl.url, detail: imageUrl.detail ?? "auto" };
+      }
+      return part;
+    });
+  }
+  return [{ type: "input_text", text: String(content) }];
+}
+
+/** Extract text from content (string or array). */
+function textOf(content: unknown): string | null {
+  if (typeof content === "string") return content;
+  if (Array.isArray(content)) {
+    const texts = content
+      .filter((c: Record<string, unknown>) => c.type === "text" || c.type === "input_text")
+      .map((c: Record<string, unknown>) => c.text as string);
+    return texts.length > 0 ? texts.join("\n") : null;
+  }
+  return null;
+}
+
+/** Extract instructions from system/developer messages already in input[]. */
+function extractInstructionsFromInput(parsed: Record<string, unknown>): string | null {
+  const input = parsed.input;
+  if (!Array.isArray(input)) return null;
+
+  for (let i = 0; i < input.length; i++) {
+    const item = input[i] as Record<string, unknown>;
+    if (item.role === "system" || item.role === "developer") {
+      const text = textOf(item.content);
+      if (text) {
+        input.splice(i, 1);
+        return text;
+      }
+    }
+  }
+  return null;
+}
+
+/** Strip `id` fields from input items — Codex backend rejects them. */
+function stripInputIds(items: Record<string, unknown>[]): void {
+  for (const item of items) {
+    if ("id" in item) {
+      delete item.id;
+    }
+  }
+}
+
 /** Extract chatgpt_account_id from JWT, falling back to stored credentials. */
 function getAccountId(accessToken: string, account: number): string | undefined {
   const creds = store.get("codex", account);

package/src/routing/affinity.ts

@@ -5,7 +5,7 @@
 
 import type { QuotaPool } from "./cooldown.ts";
 
-interface AffinityEntry {
+export interface AffinityEntry {
   pool: QuotaPool;
   account: number;
   assignedAt: number;
@@ -16,84 +16,106 @@ 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>();
-const counts = new Map<string, number>();
+export class AffinityStore {
+  private map = new Map<string, AffinityEntry>();
+  private counts = new Map<string, number>();
+  private cleanupTimer: Timer | null = null;
 
-function key(threadId: string, ampProvider: string): string {
-  return `${threadId}\0${ampProvider}`;
-}
+  private key(threadId: string, ampProvider: string): string {
+    return `${threadId}\0${ampProvider}`;
+  }
 
-function countKey(pool: QuotaPool, account: number): string {
-  return `${pool}:${account}`;
-}
+  private countKey(pool: QuotaPool, account: number): string {
+    return `${pool}:${account}`;
+  }
 
-function incCount(pool: QuotaPool, account: number): void {
-  const k = countKey(pool, account);
-  counts.set(k, (counts.get(k) ?? 0) + 1);
-}
+  private incCount(pool: QuotaPool, account: number): void {
+    const k = this.countKey(pool, account);
+    this.counts.set(k, (this.counts.get(k) ?? 0) + 1);
+  }
 
-function decCount(pool: QuotaPool, account: number): void {
-  const k = countKey(pool, account);
-  const v = (counts.get(k) ?? 0) - 1;
-  if (v <= 0) counts.delete(k);
-  else counts.set(k, v);
-}
+  private decCount(pool: QuotaPool, account: number): void {
+    const k = this.countKey(pool, account);
+    const v = (this.counts.get(k) ?? 0) - 1;
+    if (v <= 0) this.counts.delete(k);
+    else this.counts.set(k, v);
+  }
 
-export function get(threadId: string, ampProvider: string): AffinityEntry | undefined {
-  const k = key(threadId, ampProvider);
-  const entry = map.get(k);
-  if (!entry) return undefined;
-  if (Date.now() - entry.assignedAt > TTL_MS) {
-    map.delete(k);
-    decCount(entry.pool, entry.account);
-    return undefined;
+  private removeExpired(k: string, entry: AffinityEntry): void {
+    this.map.delete(k);
+    this.decCount(entry.pool, entry.account);
   }
-  // Touch: keep affinity alive while thread is active
-  entry.assignedAt = Date.now();
-  return entry;
-}
 
-export function set(threadId: string, ampProvider: string, pool: QuotaPool, account: number): void {
-  const k = key(threadId, ampProvider);
-  const existing = map.get(k);
-  if (existing) {
-    if (existing.pool !== pool || existing.account !== account) {
-      decCount(existing.pool, existing.account);
-      incCount(pool, account);
+  /** Read affinity without side effects. Returns undefined if expired or missing. */
+  peek(threadId: string, ampProvider: string): AffinityEntry | undefined {
+    const k = this.key(threadId, ampProvider);
+    const entry = this.map.get(k);
+    if (!entry) return undefined;
+    if (Date.now() - entry.assignedAt > TTL_MS) {
+      this.removeExpired(k, entry);
+      return undefined;
     }
-  } else {
-    incCount(pool, account);
+    return entry;
   }
-  map.set(k, { pool, account, assignedAt: Date.now() });
-}
 
-/** Break affinity when account is exhausted — allow re-routing. */
-export function clear(threadId: string, ampProvider: string): void {
-  const k = key(threadId, ampProvider);
-  const existing = map.get(k);
-  if (existing) {
-    decCount(existing.pool, existing.account);
-    map.delete(k);
+  /** 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();
+    return entry;
   }
-}
 
-/** Count active threads pinned to a specific (pool, account). */
-export function activeCount(pool: QuotaPool, account: number): number {
-  return counts.get(countKey(pool, account)) ?? 0;
-}
+  set(threadId: string, ampProvider: string, pool: QuotaPool, account: number): void {
+    const k = this.key(threadId, ampProvider);
+    const existing = this.map.get(k);
+    if (existing) {
+      if (existing.pool !== pool || existing.account !== account) {
+        this.decCount(existing.pool, existing.account);
+        this.incCount(pool, account);
+      }
+    } else {
+      this.incCount(pool, account);
+    }
+    this.map.set(k, { pool, account, assignedAt: Date.now() });
+  }
+
+  /** Break affinity when account is exhausted — allow re-routing. */
+  clear(threadId: string, ampProvider: string): void {
+    const k = this.key(threadId, ampProvider);
+    const existing = this.map.get(k);
+    if (existing) {
+      this.decCount(existing.pool, existing.account);
+      this.map.delete(k);
+    }
+  }
 
-let _cleanupTimer: Timer | null = null;
-
-/** Start periodic cleanup of expired entries. Call once at server startup. */
-export function startCleanup(): void {
-  if (_cleanupTimer) return;
-  _cleanupTimer = setInterval(() => {
-    const now = Date.now();
-    for (const [k, entry] of map) {
-      if (now - entry.assignedAt > TTL_MS) {
-        map.delete(k);
-        decCount(entry.pool, entry.account);
+  /** Count active threads pinned to a specific (pool, account). */
+  activeCount(pool: QuotaPool, account: number): number {
+    return this.counts.get(this.countKey(pool, account)) ?? 0;
+  }
+
+  /** Start periodic cleanup of expired entries. Call once at server startup. */
+  startCleanup(): void {
+    if (this.cleanupTimer) return;
+    this.cleanupTimer = setInterval(() => {
+      const now = Date.now();
+      for (const [k, entry] of this.map) {
+        if (now - entry.assignedAt > TTL_MS) {
+          this.removeExpired(k, entry);
+        }
       }
+    }, CLEANUP_INTERVAL_MS);
+  }
+
+  reset(): void {
+    this.map.clear();
+    this.counts.clear();
+    if (this.cleanupTimer) {
+      clearInterval(this.cleanupTimer);
+      this.cleanupTimer = null;
     }
-  }, CLEANUP_INTERVAL_MS);
+  }
 }
+
+/** Singleton instance for production use. */
+export const affinity = new AffinityStore();

package/src/routing/cooldown.ts

@@ -20,63 +20,63 @@ const EXHAUSTED_CONSECUTIVE = 3;
 /** Default burst cooldown when no Retry-After header. */
 const DEFAULT_BURST_S = 30;
 
-const entries = new Map<string, CooldownEntry>();
+export class CooldownTracker {
+  private entries = new Map<string, CooldownEntry>();
 
-function key(pool: QuotaPool, account: number): string {
-  return `${pool}:${account}`;
-}
+  private key(pool: QuotaPool, account: number): string {
+    return `${pool}:${account}`;
+  }
 
-export function isCoolingDown(pool: QuotaPool, account: number): boolean {
-  const k = key(pool, account);
-  const entry = entries.get(k);
-  if (!entry) return false;
-  if (Date.now() >= entry.until) {
-    entries.delete(k);
-    return false;
+  private getEntry(pool: QuotaPool, account: number): CooldownEntry | undefined {
+    const k = this.key(pool, account);
+    const entry = this.entries.get(k);
+    if (!entry) return undefined;
+    if (Date.now() >= entry.until) {
+      this.entries.delete(k);
+      return undefined;
+    }
+    return entry;
   }
-  return true;
-}
 
-export function isExhausted(pool: QuotaPool, account: number): boolean {
-  const k = key(pool, account);
-  const entry = entries.get(k);
-  if (!entry) return false;
-  if (Date.now() >= entry.until) {
-    entries.delete(k);
-    return false;
+  isCoolingDown(pool: QuotaPool, account: number): boolean {
+    return this.getEntry(pool, account) !== undefined;
   }
-  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 };
+  isExhausted(pool: QuotaPool, account: number): boolean {
+    return this.getEntry(pool, account)?.exhausted ?? false;
+  }
+
+  record429(pool: QuotaPool, account: number, retryAfterSeconds?: number): void {
+    const k = this.key(pool, account);
+    const entry = this.entries.get(k) ?? { until: 0, exhausted: false, consecutive429: 0 };
+
+    entry.consecutive429++;
+    const retryAfter = retryAfterSeconds ?? DEFAULT_BURST_S;
 
-  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 });
+    }
 
-  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 });
+    this.entries.set(k, entry);
   }
 
-  entries.set(k, entry);
-}
+  recordSuccess(pool: QuotaPool, account: number): void {
+    this.entries.delete(this.key(pool, account));
+  }
 
-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);
+  reset(): void {
+    this.entries.clear();
   }
 }
 
+/** Singleton instance for production use. */
+export const cooldown = new CooldownTracker();
+
 /** Parse Retry-After header (seconds or HTTP-date). */
 export function parseRetryAfter(header: string | null): number | undefined {
   if (!header) return undefined;

package/src/routing/retry.ts

@@ -3,7 +3,7 @@
 import type { ProxyConfig } from "../config/config.ts";
 import type { ParsedBody } from "../server/body.ts";
 import { logger } from "../utils/logger.ts";
-import { parseRetryAfter, record429 } from "./cooldown.ts";
+import { cooldown, parseRetryAfter } from "./cooldown.ts";
 import { type RouteResult, recordSuccess, rerouteAfter429 } from "./router.ts";
 
 /** Max 429-reroute attempts before falling back to upstream. */
@@ -33,7 +33,7 @@ export async function tryWithCachePreserve(
   }
   if (response.status === 429) {
     const nextRetryAfter = parseRetryAfter(response.headers.get("retry-after"));
-    record429(route.pool!, route.account, nextRetryAfter);
+    cooldown.record429(route.pool!, route.account, nextRetryAfter);
   }
   return null;
 }
@@ -66,7 +66,7 @@ export async function tryReroute(
 
     if (response.status === 429 && next.pool) {
       const nextRetryAfter = parseRetryAfter(response.headers.get("retry-after"));
-      record429(next.pool, next.account, nextRetryAfter);
+      cooldown.record429(next.pool, next.account, nextRetryAfter);
       currentPool = next.pool;
       currentAccount = next.account;
       continue;

package/src/routing/router.ts

@@ -17,9 +17,8 @@ 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";
+import { affinity } from "./affinity.ts";
+import { cooldown, type QuotaPool } from "./cooldown.ts";
 
 interface ProviderEntry {
   provider: Provider;

package/src/server/server.ts

@@ -1,15 +1,16 @@
 /** HTTP server — routes provider requests through local OAuth or Amp upstream. */
 
+import { maybeShowAd } from "../cli/ads.ts";
 import type { ProxyConfig } from "../config/config.ts";
 import * as rewriter from "../proxy/rewriter.ts";
 import * as upstream from "../proxy/upstream.ts";
-import { startCleanup } from "../routing/affinity.ts";
+import { affinity } from "../routing/affinity.ts";
 import { tryReroute, tryWithCachePreserve } from "../routing/retry.ts";
 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 { record, snapshot } from "../utils/stats.ts";
+import { stats } from "../utils/stats.ts";
 import { type ParsedBody, parseBody } from "./body.ts";
 
 export function startServer(config: ProxyConfig): ReturnType<typeof Bun.serve> {
@@ -34,7 +35,7 @@ export function startServer(config: ProxyConfig): ReturnType<typeof Bun.serve> {
     },
   });
 
-  startCleanup();
+  affinity.startCleanup();
   logger.info(`ampcode-connector listening on http://localhost:${config.port}`);
 
   const shutdown = () => {
@@ -129,7 +130,7 @@ async function handleProvider(
     response = await fallbackUpstream(req, body, config);
   }
 
-  record({
+  stats.record({
     timestamp: new Date().toISOString(),
     route: route.decision,
     provider: providerName,
@@ -138,6 +139,8 @@ async function handleProvider(
     durationMs: Date.now() - startTime,
   });
 
+  maybeShowAd();
+
   return response;
 }
 
@@ -158,6 +161,6 @@ function healthCheck(config: ProxyConfig): Response {
     port: config.port,
     upstream: config.ampUpstreamUrl,
     providers: config.providers,
-    stats: snapshot(),
+    stats: stats.snapshot(),
   });
 }

package/src/utils/stats.ts

@@ -11,22 +11,6 @@ export interface RequestEntry {
   durationMs: number;
 }
 
-const MAX_ENTRIES = 1000;
-const buffer: RequestEntry[] = [];
-let writeIndex = 0;
-let totalCount = 0;
-const startedAt = Date.now();
-
-export function record(entry: RequestEntry): void {
-  if (buffer.length < MAX_ENTRIES) {
-    buffer.push(entry);
-  } else {
-    buffer[writeIndex] = entry;
-  }
-  writeIndex = (writeIndex + 1) % MAX_ENTRIES;
-  totalCount++;
-}
-
 export interface StatsSnapshot {
   totalRequests: number;
   requestsByRoute: Partial<Record<RouteDecision, number>>;
@@ -35,35 +19,66 @@ export interface StatsSnapshot {
   uptimeMs: number;
 }
 
-export function snapshot(): StatsSnapshot {
-  const requestsByRoute: Partial<Record<RouteDecision, number>> = {};
-  let count429 = 0;
-  let totalDuration = 0;
+export class StatsRecorder {
+  private readonly maxEntries: number;
+  private buffer: RequestEntry[] = [];
+  private writeIndex = 0;
+  private totalCount = 0;
+  private readonly startedAt = Date.now();
 
-  for (const entry of buffer) {
-    requestsByRoute[entry.route] = (requestsByRoute[entry.route] ?? 0) + 1;
-    if (entry.statusCode === 429) count429++;
-    totalDuration += entry.durationMs;
+  constructor(maxEntries = 1000) {
+    this.maxEntries = maxEntries;
   }
 
-  return {
-    totalRequests: totalCount,
-    requestsByRoute,
-    count429,
-    averageDurationMs: buffer.length > 0 ? totalDuration / buffer.length : 0,
-    uptimeMs: Date.now() - startedAt,
-  };
-}
+  record(entry: RequestEntry): void {
+    if (this.buffer.length < this.maxEntries) {
+      this.buffer.push(entry);
+    } else {
+      this.buffer[this.writeIndex] = entry;
+    }
+    this.writeIndex = (this.writeIndex + 1) % this.maxEntries;
+    this.totalCount++;
+  }
+
+  snapshot(): StatsSnapshot {
+    const requestsByRoute: Partial<Record<RouteDecision, number>> = {};
+    let count429 = 0;
+    let totalDuration = 0;
 
-export function recentRequests(n: number): RequestEntry[] {
-  const count = Math.min(n, buffer.length);
-  if (count === 0) return [];
+    for (const entry of this.buffer) {
+      requestsByRoute[entry.route] = (requestsByRoute[entry.route] ?? 0) + 1;
+      if (entry.statusCode === 429) count429++;
+      totalDuration += entry.durationMs;
+    }
 
-  const result: RequestEntry[] = [];
-  let idx = (writeIndex - count + buffer.length) % buffer.length;
-  for (let i = 0; i < count; i++) {
-    result.push(buffer[idx]!);
-    idx = (idx + 1) % buffer.length;
+    return {
+      totalRequests: this.totalCount,
+      requestsByRoute,
+      count429,
+      averageDurationMs: this.buffer.length > 0 ? totalDuration / this.buffer.length : 0,
+      uptimeMs: Date.now() - this.startedAt,
+    };
+  }
+
+  recentRequests(n: number): RequestEntry[] {
+    const count = Math.min(n, this.buffer.length);
+    if (count === 0) return [];
+
+    const result: RequestEntry[] = [];
+    let idx = (this.writeIndex - count + this.buffer.length) % this.buffer.length;
+    for (let i = 0; i < count; i++) {
+      result.push(this.buffer[idx]!);
+      idx = (idx + 1) % this.buffer.length;
+    }
+    return result;
+  }
+
+  reset(): void {
+    this.buffer = [];
+    this.writeIndex = 0;
+    this.totalCount = 0;
   }
-  return result;
 }
+
+/** Singleton instance for production use. */
+export const stats = new StatsRecorder();