ampcode-connector 0.1.4 → 0.1.5

package/package.json

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

package/src/auth/auto-refresh.ts

@@ -0,0 +1,44 @@
+import { TOKEN_EXPIRY_BUFFER_MS } from "../constants.ts";
+import { logger } from "../utils/logger.ts";
+import * as configs from "./configs.ts";
+import type { OAuthConfig } from "./oauth.ts";
+import * as oauth from "./oauth.ts";
+import { getAll, type ProviderName } from "./store.ts";
+
+const REFRESH_INTERVAL_MS = 60_000;
+
+const providerConfigs: Record<ProviderName, OAuthConfig> = {
+  anthropic: configs.anthropic,
+  codex: configs.codex,
+  google: configs.google,
+};
+
+let timer: Timer | null = null;
+
+async function refreshAll(): Promise<void> {
+  const now = Date.now();
+
+  for (const [provider, config] of Object.entries(providerConfigs) as [ProviderName, OAuthConfig][]) {
+    for (const { account, credentials } of getAll(provider)) {
+      if (credentials.expiresAt - now > TOKEN_EXPIRY_BUFFER_MS) continue;
+
+      try {
+        logger.debug("Auto-refreshing token", { provider, account });
+        await oauth.token(config, account);
+      } catch (err) {
+        logger.error("Auto-refresh failed", { provider, account, error: String(err) });
+      }
+    }
+  }
+}
+
+export function startAutoRefresh(): void {
+  if (timer) return;
+  timer = setInterval(refreshAll, REFRESH_INTERVAL_MS);
+}
+
+export function stopAutoRefresh(): void {
+  if (!timer) return;
+  clearInterval(timer);
+  timer = null;
+}

package/src/auth/callback-server.ts

@@ -17,7 +17,7 @@ export async function waitForCallback(
 ): Promise<CallbackResult> {
   return new Promise((resolve, reject) => {
     let server: ReturnType<typeof Bun.serve> | null = null;
-    let timeoutId: ReturnType<typeof setTimeout> | null = null;
+    let timeoutId: Timer | null = null;
 
     const cleanup = () => {
       if (timeoutId) clearTimeout(timeoutId);

package/src/auth/discovery.ts

@@ -41,7 +41,7 @@ function accountIdFromJWT(token: string): string | null {
     if (parts.length < 2 || !parts[1]) return null;
     const payload = JSON.parse(new TextDecoder().decode(fromBase64url(parts[1]))) as Raw;
     const auth = payload["https://api.openai.com/auth"] as Raw | undefined;
-    return (auth?.["chatgpt_account_id"] as string) ?? null;
+    return (auth?.chatgpt_account_id as string) ?? null;
   } catch {
     return null;
   }

package/src/auth/oauth.ts

@@ -39,8 +39,14 @@ export async function token(config: OAuthConfig, account = 0): Promise<string |
   try {
     return (await refresh(config, creds.refreshToken, account)).accessToken;
   } catch (err) {
-    logger.error(`Token refresh failed for ${config.providerName}:${account}`, { error: String(err) });
-    return null;
+    logger.warn(`Token refresh failed for ${config.providerName}:${account}, retrying...`, { error: String(err) });
+    try {
+      await Bun.sleep(1000);
+      return (await refresh(config, creds.refreshToken, account)).accessToken;
+    } catch (retryErr) {
+      logger.error(`Token refresh retry failed for ${config.providerName}:${account}`, { error: String(retryErr) });
+      return null;
+    }
   }
 }
 
@@ -54,9 +60,7 @@ export async function tokenFromAny(config: OAuthConfig): Promise<{ accessToken:
     try {
       const refreshed = await refresh(config, c.refreshToken, account);
       return { accessToken: refreshed.accessToken, account };
-    } catch {
-      continue;
-    }
+    } catch {}
   }
 
   return null;
@@ -102,7 +106,7 @@ export async function login(config: OAuthConfig): Promise<Credentials> {
     redirect_uri: redirectUri,
     code_verifier: verifier,
   };
-  if (config.sendStateInExchange) exchangeParams["state"] = callback.state;
+  if (config.sendStateInExchange) exchangeParams.state = callback.state;
 
   const raw = await exchange(config, exchangeParams);
   let credentials = parseTokenFields(raw, config);
@@ -156,7 +160,7 @@ async function refresh(config: OAuthConfig, refreshToken: string, account = 0):
 
 async function exchange(config: OAuthConfig, params: Record<string, string>): Promise<Record<string, unknown>> {
   const all: Record<string, string> = { client_id: config.clientId, ...params };
-  if (config.clientSecret) all["client_secret"] = config.clientSecret;
+  if (config.clientSecret) all.client_secret = config.clientSecret;
 
   const isJson = config.bodyFormat === "json";
   const res = await fetch(config.tokenUrl, {

package/src/auth/store.ts

@@ -2,7 +2,7 @@
  *  Multi-account: composite key (provider, account).
  *  Sync API via bun:sqlite — no cache needed at 0.4µs/read. */
 
-import { Database } from "bun:sqlite";
+import { Database, type Statement } from "bun:sqlite";
 import { mkdirSync } from "node:fs";
 import { homedir } from "node:os";
 import { join } from "node:path";
@@ -21,68 +21,100 @@ export type ProviderName = "anthropic" | "codex" | "google";
 const DIR = join(homedir(), ".ampcode-connector");
 const DB_PATH = join(DIR, "credentials.db");
 
-mkdirSync(DIR, { recursive: true, mode: 0o700 });
-
-const db = new Database(DB_PATH, { strict: true });
-db.exec("PRAGMA journal_mode=WAL");
-
-db.exec(`
-  CREATE TABLE IF NOT EXISTS credentials (
-    provider TEXT NOT NULL,
-    account  INTEGER NOT NULL DEFAULT 0,
-    data     TEXT NOT NULL,
-    PRIMARY KEY (provider, account)
-  )
-`);
-
-const stmtGet = db.prepare<{ data: string }, [string, number]>(
-  "SELECT data FROM credentials WHERE provider = ? AND account = ?",
-);
-const stmtGetAll = db.prepare<{ account: number; data: string }, [string]>(
-  "SELECT account, data FROM credentials WHERE provider = ? ORDER BY account",
-);
-const stmtSet = db.prepare<void, [string, number, string]>(
-  "INSERT OR REPLACE INTO credentials (provider, account, data) VALUES (?, ?, ?)",
-);
-const stmtDelOne = db.prepare<void, [string, number]>("DELETE FROM credentials WHERE provider = ? AND account = ?");
-const stmtDelAll = db.prepare<void, [string]>("DELETE FROM credentials WHERE provider = ?");
-const stmtMaxAccount = db.prepare<{ max_account: number | null }, [string]>(
-  "SELECT MAX(account) as max_account FROM credentials WHERE provider = ?",
-);
-const stmtCount = db.prepare<{ cnt: number }, [string]>("SELECT COUNT(*) as cnt FROM credentials WHERE provider = ?");
+interface DataRow {
+  data: string;
+}
+interface AccountDataRow {
+  account: number;
+  data: string;
+}
+interface MaxAccountRow {
+  max_account: number | null;
+}
+interface CountRow {
+  cnt: number;
+}
+
+interface Statements {
+  get: Statement<DataRow, [string, number]>;
+  getAll: Statement<AccountDataRow, [string]>;
+  set: Statement<void, [string, number, string]>;
+  delOne: Statement<void, [string, number]>;
+  delAll: Statement<void, [string]>;
+  maxAccount: Statement<MaxAccountRow, [string]>;
+  count: Statement<CountRow, [string]>;
+}
+
+let _db: Database | null = null;
+let _stmts: Statements | null = null;
+
+function init() {
+  if (_stmts) return _stmts;
+
+  mkdirSync(DIR, { recursive: true, mode: 0o700 });
+  _db = new Database(DB_PATH, { strict: true });
+  _db.exec("PRAGMA journal_mode=WAL");
+  _db.exec(`
+    CREATE TABLE IF NOT EXISTS credentials (
+      provider TEXT NOT NULL,
+      account  INTEGER NOT NULL DEFAULT 0,
+      data     TEXT NOT NULL,
+      PRIMARY KEY (provider, account)
+    )
+  `);
+
+  _stmts = {
+    get: _db.prepare<DataRow, [string, number]>("SELECT data FROM credentials WHERE provider = ? AND account = ?"),
+    getAll: _db.prepare<AccountDataRow, [string]>(
+      "SELECT account, data FROM credentials WHERE provider = ? ORDER BY account",
+    ),
+    set: _db.prepare<void, [string, number, string]>(
+      "INSERT OR REPLACE INTO credentials (provider, account, data) VALUES (?, ?, ?)",
+    ),
+    delOne: _db.prepare<void, [string, number]>("DELETE FROM credentials WHERE provider = ? AND account = ?"),
+    delAll: _db.prepare<void, [string]>("DELETE FROM credentials WHERE provider = ?"),
+    maxAccount: _db.prepare<MaxAccountRow, [string]>(
+      "SELECT MAX(account) as max_account FROM credentials WHERE provider = ?",
+    ),
+    count: _db.prepare<CountRow, [string]>("SELECT COUNT(*) as cnt FROM credentials WHERE provider = ?"),
+  };
+  return _stmts;
+}
 
 export function get(provider: ProviderName, account = 0): Credentials | undefined {
-  const row = stmtGet.get(provider, account);
+  const row = init().get.get(provider, account);
   if (!row) return undefined;
   return JSON.parse(row.data) as Credentials;
 }
 
 export function getAll(provider: ProviderName): { account: number; credentials: Credentials }[] {
-  return stmtGetAll.all(provider).map((row) => ({
-    account: row.account,
-    credentials: JSON.parse(row.data) as Credentials,
-  }));
+  return init()
+    .getAll.all(provider)
+    .map((row) => ({
+      account: row.account,
+      credentials: JSON.parse(row.data) as Credentials,
+    }));
 }
 
 export function save(provider: ProviderName, credentials: Credentials, account = 0): void {
-  stmtSet.run(provider, account, JSON.stringify(credentials));
+  init().set.run(provider, account, JSON.stringify(credentials));
 }
 
 export function remove(provider: ProviderName, account?: number): void {
   if (account !== undefined) {
-    stmtDelOne.run(provider, account);
+    init().delOne.run(provider, account);
   } else {
-    stmtDelAll.run(provider);
+    init().delAll.run(provider);
   }
 }
 
 export function nextAccount(provider: ProviderName): number {
-  const row = stmtMaxAccount.get(provider);
+  const row = init().maxAccount.get(provider);
   return (row?.max_account ?? -1) + 1;
 }
 
 export function count(provider: ProviderName): number {
-  return stmtCount.get(provider)?.cnt ?? 0;
+  return init().count.get(provider)?.cnt ?? 0;
 }
 
 /** Find existing account by email or accountId match. */

package/src/cli/setup.ts

@@ -16,7 +16,7 @@ const AMP_SETTINGS_PATHS = [
 ];
 
 function ampSettingsPaths(): string[] {
-  const envPath = process.env["AMP_SETTINGS_FILE"];
+  const envPath = process.env.AMP_SETTINGS_FILE;
   return envPath ? [envPath] : AMP_SETTINGS_PATHS;
 }
 
@@ -31,11 +31,11 @@ function readJson(path: string): Record<string, unknown> {
 
 function writeJson(path: string, data: Record<string, unknown>): void {
   mkdirSync(dirname(path), { recursive: true });
-  writeFileSync(path, JSON.stringify(data, null, 2) + "\n", "utf-8");
+  writeFileSync(path, `${JSON.stringify(data, null, 2)}\n`, "utf-8");
 }
 
 function findAmpApiKey(proxyUrl: string): string | undefined {
-  if (process.env["AMP_API_KEY"]) return process.env["AMP_API_KEY"];
+  if (process.env.AMP_API_KEY) return process.env.AMP_API_KEY;
 
   const secrets = readJson(AMP_SECRETS_PATH);
   const exact = secrets[`apiKey@${proxyUrl}`];
@@ -102,7 +102,7 @@ export async function setup(): Promise<void> {
 
   if (existingKey) {
     saveAmpApiKey(existingKey, proxyUrl);
-    const preview = existingKey.slice(0, 10) + "...";
+    const preview = `${existingKey.slice(0, 10)}...`;
     line(`${s.green}ok${s.reset} Amp token found  ${s.dim}${preview}${s.reset}`);
   } else {
     line();

package/src/cli/tui.ts

@@ -27,7 +27,7 @@ let selected = 0;
 let items: Item[] = [];
 let message = "";
 let busy = false;
-let timer: ReturnType<typeof setInterval> | null = null;
+let timer: Timer | null = null;
 
 export function dashboard(): void {
   if (!process.stdin.isTTY) throw new Error("Interactive dashboard requires a TTY.");

package/src/config/config.ts

@@ -36,18 +36,18 @@ const SECRETS_PATH = join(homedir(), ".local", "share", "amp", "secrets.json");
 export async function loadConfig(): Promise<ProxyConfig> {
   const file = await readConfigFile();
   const apiKey = await resolveApiKey(file);
-  const providers = asRecord(file?.["providers"]);
+  const providers = asRecord(file?.providers);
 
   return {
-    port: asNumber(file?.["port"]) ?? DEFAULTS.port,
-    ampUpstreamUrl: asString(file?.["ampUpstreamUrl"]) ?? DEFAULTS.ampUpstreamUrl,
+    port: asNumber(file?.port) ?? DEFAULTS.port,
+    ampUpstreamUrl: asString(file?.ampUpstreamUrl) ?? DEFAULTS.ampUpstreamUrl,
     ampApiKey: apiKey,
-    exaApiKey: asString(file?.["exaApiKey"]) ?? process.env["EXA_API_KEY"],
-    logLevel: asLogLevel(file?.["logLevel"]) ?? DEFAULTS.logLevel,
+    exaApiKey: asString(file?.exaApiKey) ?? process.env.EXA_API_KEY,
+    logLevel: asLogLevel(file?.logLevel) ?? DEFAULTS.logLevel,
     providers: {
-      anthropic: asBool(providers?.["anthropic"]) ?? DEFAULTS.providers.anthropic,
-      codex: asBool(providers?.["codex"]) ?? DEFAULTS.providers.codex,
-      google: asBool(providers?.["google"]) ?? DEFAULTS.providers.google,
+      anthropic: asBool(providers?.anthropic) ?? DEFAULTS.providers.anthropic,
+      codex: asBool(providers?.codex) ?? DEFAULTS.providers.codex,
+      google: asBool(providers?.google) ?? DEFAULTS.providers.google,
     },
   };
 }
@@ -70,10 +70,10 @@ async function readConfigFile(): Promise<Record<string, unknown> | null> {
 
 /** Amp API key resolution: config file → AMP_API_KEY env → secrets.json */
 async function resolveApiKey(file: Record<string, unknown> | null): Promise<string | undefined> {
-  const fromFile = asString(file?.["ampApiKey"]);
+  const fromFile = asString(file?.ampApiKey);
   if (fromFile) return fromFile;
 
-  const fromEnv = process.env["AMP_API_KEY"];
+  const fromEnv = process.env.AMP_API_KEY;
   if (fromEnv) return fromEnv;
 
   return readSecretsFile();
@@ -101,7 +101,7 @@ function asRecord(v: unknown): Record<string, unknown> | undefined {
 }
 
 function asNumber(v: unknown): number | undefined {
-  return typeof v === "number" && !isNaN(v) ? v : undefined;
+  return typeof v === "number" && !Number.isNaN(v) ? v : undefined;
 }
 
 function asString(v: unknown): string | undefined {

package/src/index.ts

@@ -1,13 +1,15 @@
 #!/usr/bin/env bun
 /** ampcode-connector entry point. */
 
+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 { line, s } from "./cli/ansi.ts";
 import { setup } from "./cli/setup.ts";
+import * as status from "./cli/status.ts";
 import { dashboard } from "./cli/tui.ts";
-import { loadConfig } from "./config/config.ts";
+import { loadConfig, type ProxyConfig } from "./config/config.ts";
 import { startServer } from "./server/server.ts";
 import { logger, setLogLevel } from "./utils/logger.ts";
 
@@ -42,12 +44,36 @@ async function main(): Promise<void> {
   const config = await loadConfig();
   setLogLevel(config.logLevel);
   startServer(config);
+  startAutoRefresh();
+  banner(config);
 
   // Non-blocking update check — runs in background after server starts
   const { checkForUpdates } = await import("./utils/update-check.ts");
   checkForUpdates();
 }
 
+function banner(config: ProxyConfig): void {
+  const providers = status.all();
+  const upstream = config.ampUpstreamUrl.replace(/^https?:\/\//, "");
+
+  line();
+  line(`  ${s.bold}ampcode-connector${s.reset}`);
+  line(`  ${s.dim}http://localhost:${config.port}${s.reset}`);
+  line();
+
+  for (const p of providers) {
+    const count = p.accounts.length;
+    const label = p.label.padEnd(16);
+    const countStr = count > 0 ? `${count} account${count > 1 ? "s" : ""}` : "--";
+    const dot = count > 0 ? `${s.green}●${s.reset}` : `${s.dim}○${s.reset}`;
+    line(`  ${label}  ${countStr.padEnd(12)}${dot}`);
+  }
+
+  line();
+  line(`  ${s.dim}upstream → ${upstream}${s.reset}`);
+  line();
+}
+
 function usage(): void {
   line();
   line(`${s.bold}ampcode-connector${s.reset} ${s.dim}— proxy Amp CLI through local OAuth subscriptions${s.reset}`);

package/src/providers/antigravity.ts

@@ -86,8 +86,5 @@ async function tryEndpoints(
     }
   }
 
-  return new Response(JSON.stringify({ error: `All Antigravity endpoints failed: ${lastError?.message}` }), {
-    status: 502,
-    headers: { "Content-Type": "application/json" },
-  });
+  return Response.json({ error: `All Antigravity endpoints failed: ${lastError?.message}` }, { status: 502 });
 }

package/src/providers/base.ts

@@ -55,8 +55,5 @@ export async function forward(opts: ForwardOptions): Promise<Response> {
 }
 
 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" },
-  });
+  return Response.json({ error: `No ${providerName} OAuth token available. Run login first.` }, { status: 401 });
 }

package/src/providers/codex.ts

@@ -58,7 +58,7 @@ function getAccountId(accessToken: string, account: number): string | undefined
     if (parts.length < 2 || !parts[1]) return undefined;
     const payload = JSON.parse(new TextDecoder().decode(fromBase64url(parts[1]))) as Record<string, unknown>;
     const auth = payload["https://api.openai.com/auth"] as Record<string, unknown> | undefined;
-    return (auth?.["chatgpt_account_id"] as string) ?? undefined;
+    return (auth?.chatgpt_account_id as string) ?? undefined;
   } catch {
     return undefined;
   }

package/src/proxy/rewriter.ts

@@ -7,6 +7,9 @@
 
 import { modelFieldPaths } from "../constants.ts";
 
+/** Pre-split field paths — avoids .split(".") on every SSE chunk. */
+const MODEL_FIELD_PARTS = modelFieldPaths.map((p) => p.split("."));
+
 export function rewrite(originalModel: string): (data: string) => string {
   return (data: string) => {
     if (data === "[DONE]") return data;
@@ -15,10 +18,10 @@ export function rewrite(originalModel: string): (data: string) => string {
       const parsed = JSON.parse(data) as Record<string, unknown>;
       let modified = false;
 
-      for (const path of modelFieldPaths) {
-        const current = getField(parsed, path);
+      for (const parts of MODEL_FIELD_PARTS) {
+        const current = getField(parsed, parts);
         if (current !== undefined && current !== originalModel) {
-          setField(parsed, path, originalModel);
+          setField(parsed, parts, originalModel);
           modified = true;
         }
       }
@@ -32,8 +35,7 @@ export function rewrite(originalModel: string): (data: string) => string {
   };
 }
 
-function getField(obj: Record<string, unknown>, path: string): unknown {
-  const parts = path.split(".");
+function getField(obj: Record<string, unknown>, parts: string[]): unknown {
   let current: unknown = obj;
   for (const part of parts) {
     if (current == null || typeof current !== "object") return undefined;
@@ -42,8 +44,7 @@ function getField(obj: Record<string, unknown>, path: string): unknown {
   return current;
 }
 
-function setField(obj: Record<string, unknown>, path: string, value: unknown): void {
-  const parts = path.split(".");
+function setField(obj: Record<string, unknown>, parts: string[], value: unknown): void {
   let current: unknown = obj;
   for (let i = 0; i < parts.length - 1; i++) {
     if (current == null || typeof current !== "object") return;
@@ -55,15 +56,15 @@ function setField(obj: Record<string, unknown>, path: string, value: unknown): v
 }
 
 function suppressThinking(data: Record<string, unknown>): boolean {
-  const content = data["content"];
+  const content = data.content;
   if (!Array.isArray(content)) return false;
 
-  const hasToolUse = content.some((b: Record<string, unknown>) => b["type"] === "tool_use");
+  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");
+  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");
+  data.content = content.filter((b: Record<string, unknown>) => b.type !== "thinking");
   return true;
 }

package/src/proxy/upstream.ts

@@ -33,9 +33,6 @@ export async function forward(request: Request, ampUpstreamUrl: string, ampApiKe
     });
   } 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" },
-    });
+    return Response.json({ error: "Failed to connect to Amp upstream", details: String(err) }, { status: 502 });
   }
 }

package/src/routing/affinity.ts

@@ -17,16 +17,35 @@ const TTL_MS = 2 * 3600_000;
 const CLEANUP_INTERVAL_MS = 10 * 60_000;
 
 const map = new Map<string, AffinityEntry>();
+const counts = new Map<string, number>();
 
 function key(threadId: string, ampProvider: string): string {
   return `${threadId}\0${ampProvider}`;
 }
 
+function 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);
+}
+
+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);
+}
+
 export function get(threadId: string, ampProvider: string): AffinityEntry | undefined {
-  const entry = map.get(key(threadId, ampProvider));
+  const k = key(threadId, ampProvider);
+  const entry = map.get(k);
   if (!entry) return undefined;
   if (Date.now() - entry.assignedAt > TTL_MS) {
-    map.delete(key(threadId, ampProvider));
+    map.delete(k);
+    decCount(entry.pool, entry.account);
     return undefined;
   }
   // Touch: keep affinity alive while thread is active
@@ -35,30 +54,46 @@ export function get(threadId: string, ampProvider: string): AffinityEntry | unde
 }
 
 export function set(threadId: string, ampProvider: string, pool: QuotaPool, account: number): void {
-  map.set(key(threadId, ampProvider), { pool, account, assignedAt: Date.now() });
+  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);
+    }
+  } else {
+    incCount(pool, account);
+  }
+  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 {
-  map.delete(key(threadId, ampProvider));
+  const k = key(threadId, ampProvider);
+  const existing = map.get(k);
+  if (existing) {
+    decCount(existing.pool, existing.account);
+    map.delete(k);
+  }
 }
 
 /** 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;
+  return counts.get(countKey(pool, account)) ?? 0;
 }
 
-/** Periodic cleanup of expired entries. */
-setInterval(() => {
-  const now = Date.now();
-  for (const [k, entry] of map) {
-    if (now - entry.assignedAt > TTL_MS) map.delete(k);
-  }
-}, CLEANUP_INTERVAL_MS);
+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);
+      }
+    }
+  }, CLEANUP_INTERVAL_MS);
+}

package/src/routing/cooldown.ts

@@ -11,8 +11,6 @@ interface CooldownEntry {
   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. */
@@ -29,20 +27,22 @@ function key(pool: QuotaPool, account: number): string {
 }
 
 export function isCoolingDown(pool: QuotaPool, account: number): boolean {
-  const entry = entries.get(key(pool, account));
+  const k = key(pool, account);
+  const entry = entries.get(k);
   if (!entry) return false;
   if (Date.now() >= entry.until) {
-    entries.delete(key(pool, account));
+    entries.delete(k);
     return false;
   }
   return true;
 }
 
 export function isExhausted(pool: QuotaPool, account: number): boolean {
-  const entry = entries.get(key(pool, account));
+  const k = key(pool, account);
+  const entry = entries.get(k);
   if (!entry) return false;
   if (Date.now() >= entry.until) {
-    entries.delete(key(pool, account));
+    entries.delete(k);
     return false;
   }
   return entry.exhausted;