maxpool 1.3.1 → 1.5.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "maxpool",
-  "version": "1.3.1",
+  "version": "1.5.0",
   "description": "Multi-account Claude Code proxy with adaptive, rate-aware load balancing across Claude accounts",
   "type": "module",
   "main": "src/index.js",
@@ -13,7 +13,7 @@
   ],
   "scripts": {
     "start": "node src/index.js",
-    "test": "node --test",
+    "test": "bash scripts/run-tests.sh",
     "lint": "eslint src/ test/",
     "release": "bash scripts/release.sh"
   },

package/src/account-manager.js

@@ -9,6 +9,16 @@ import { refreshAccessToken, isTokenExpiringSoon } from './oauth.js';
 // an Infinity session-kill / error-fast.
 const BOUNDED_REPOLL_HOLD_MS = 60_000;
 
+// Session rebalancing (issue #1): a bound session may migrate OFF a hot account
+// onto a much-healthier one, but ONLY on a thinking-safe request (see
+// _migrationSafeForRequest). These margins force a CLEAR, flap-stable win so a
+// session can never ping-pong between two similarly-loaded accounts.
+const REBALANCE_SCORE_MARGIN = 0.5;   // candidate score must be ≤ 50% of the bound account's
+const REBALANCE_MIN_ABS_GAP = 0.5;    // …with a small absolute floor so near-zero scores don't micro-churn
+// Weekly-pressure tiers, healthiest first. A fresh (unknown) account is the best
+// migration target; migration requires the candidate be a STRICTLY healthier tier.
+const WEEKLY_TIER = { unknown: 0, normal: 0, soft: 1, reserve: 2, critical: 3, exhausted: 4 };
+
 function emptyQuota() {
   return {
     // Standard API rate limits (API key accounts)
@@ -191,6 +201,21 @@ export class AccountManager {
       bytes: 0, // aggregate buffered body bytes across all held requests
     };
     this.admissionPaused = false;
+    // Single-writer baton: only the lease holder may rotate OAuth refresh tokens
+    // (refresh tokens are single-use; two refreshers brick the account). A worker
+    // booted headless during a reload starts WITHOUT the lease and refreshes
+    // nothing until it acquires the baton. Default true so the standalone /
+    // direct-listen (non-supervised, headless service) path is unchanged.
+    this.writerLease = true;
+  }
+
+  /**
+   * Acquire/release the single-writer baton. While released, ensureTokenFresh is
+   * a no-op (the worker serves on its existing access tokens but never rotates a
+   * single-use refresh token — that's the lease holder's job).
+   */
+  setWriterLease(held) {
+    this.writerLease = Boolean(held);
   }
 
   /**
@@ -546,7 +571,7 @@ export class AccountManager {
     console.log(`[Maxpool] Anthropic recovery probe failed; retrying in ${retryAfter}s (${reason})`);
   }
 
-  noteAmbiguousRateLimit(accountIndex, fingerprint, retryAfterSeconds) {
+  noteAmbiguousRateLimit(accountIndex, fingerprint, _retryAfterSeconds) {
     if (!fingerprint) return false;
     const now = Date.now();
     const windowMs = 30_000;
@@ -860,6 +885,62 @@ export class AccountManager {
       || ['reserve', 'critical', 'exhausted'].includes(this._weeklyRawState(account));
   }
 
+  // ── Session rebalancing (issue #1) ────────────────────────────────────────
+  // A bound session normally sticks to its account (continuity + Anthropic signed-
+  // thinking signature validity). These let it migrate OFF a hot account onto fresh
+  // capacity, but only when it's provably safe and clearly worth it.
+
+  /** Per-request safety gate: a request is safe to migrate to a DIFFERENT account
+   *  iff its body carries NO signed thinking (replaying a signed thinking block to
+   *  another account → non-retryable "invalid signature"). Uses the PER-REQUEST
+   *  body signal only — never the session-sticky policy — and fails CLOSED on any
+   *  body we couldn't fully scan (non-JSON / parse error → bodyThinkingScanned unset). */
+  _migrationSafeForRequest(requestInfo = {}) {
+    return requestInfo.bodyThinkingScanned === true
+      && requestInfo.requiresAnthropicThinkingIntegrity !== true;
+  }
+
+  /** Flap-stable "hot": weekly reserve/critical or 5h-cap pressure — signals that
+   *  do NOT flip the instant a request migrates (unlike live in-flight, left to the
+   *  score loop). A healthy bound account is never hot, so it never migrates → no
+   *  ping-pong. */
+  _isBoundAccountHot(account) {
+    return this._isNearQuota(account);
+  }
+
+  /** Decide whether a bound session should leave its (hot) account THIS request.
+   *  All gates must hold: thinking-safe + not mid queue-admission + bound is hot +
+   *  a genuinely-healthy alternative that is BOTH much cheaper AND a strictly
+   *  healthier weekly tier (so concurrency jitter alone can never trigger a move). */
+  _shouldRebalanceBoundSession(bound, profile, excludedIndexes, requestInfo, scoringCtx) {
+    if (!this._migrationSafeForRequest(requestInfo)) return false;
+    if (requestInfo.queueTicket || requestInfo.queueAdmitted) return false;
+    if (!this._isBoundAccountHot(bound)) return false;
+
+    const boundScore = this._scoreAccount(bound, requestInfo, scoringCtx);
+    const boundTier = WEEKLY_TIER[this._weeklyRawState(bound)] ?? 0;
+
+    let bestScore = Infinity;
+    let bestTier = Infinity;
+    for (const account of this.accounts) {
+      if (account.index === bound.index) continue;
+      if (excludedIndexes.has(account.index)) continue;
+      if (!this._matchesRequest(account, profile, requestInfo)) continue;
+      // Genuinely-healthy alternatives only (normal/soft/unknown weekly).
+      if (!this._isAvailable(account, { allowWeeklyReserve: false, allowWeeklyCritical: false })) continue;
+      const score = this._scoreAccount(account, requestInfo, scoringCtx);
+      if (score < bestScore) {
+        bestScore = score;
+        bestTier = WEEKLY_TIER[this._weeklyRawState(account)] ?? 0;
+      }
+    }
+    if (!Number.isFinite(bestScore)) return false; // no healthy alternative
+
+    return bestScore <= boundScore * REBALANCE_SCORE_MARGIN
+      && (boundScore - bestScore) >= REBALANCE_MIN_ABS_GAP
+      && bestTier < boundTier;
+  }
+
   _retryInfo(account) {
     const now = Date.now();
     const q = account.quota || {};
@@ -989,7 +1070,13 @@ export class AccountManager {
       }
     }
     const bound = this._boundAccount(requestInfo.sessionKey, profile, excludedIndexes, requestInfo);
-    if (bound && !this._hasHigherPriorityAvailable(bound, profile, excludedIndexes, requestInfo)) return bound;
+    if (bound
+      && !this._hasHigherPriorityAvailable(bound, profile, excludedIndexes, requestInfo)
+      && !this._shouldRebalanceBoundSession(bound, profile, excludedIndexes, requestInfo, scoringCtx)) {
+      return bound;
+    }
+    // Else fall through to the candidate score loop, which re-homes the session
+    // onto the best healthy account via _bindSession on acquire.
 
     const weeklyPasses = hasBinding
       ? [
@@ -1106,6 +1193,16 @@ export class AccountManager {
     if (!binding.homeName || priority < binding.homePriority) {
       binding.homeName = account.name;
       binding.homePriority = priority;
+    } else if (priority === binding.homePriority && account.name !== binding.homeName) {
+      // Same-priority move. If the previous home is still AVAILABLE we left it by
+      // CHOICE (session rebalancing off a hot-but-usable account) → re-home, so
+      // _boundAccount (which prefers homeName) doesn't snap the session straight
+      // back to the hot home next request. If the old home is UNAVAILABLE this is a
+      // FAILOVER → keep homeName so the session returns to it once it recovers.
+      const oldHome = this.accounts.find(a => a.name === binding.homeName);
+      if (oldHome && this._isAvailable(oldHome, { allowWeeklyReserve: true })) {
+        binding.homeName = account.name;
+      }
     }
     binding.currentName = account.name;
     this.sessionBindings.set(sessionKey, binding);
@@ -1635,6 +1732,12 @@ export class AccountManager {
     const account = this.accounts[accountIndex];
     if (!account || account.type !== 'oauth' || !account.refreshToken) return true;
 
+    // Single-writer baton: a worker without the lease NEVER rotates a single-use
+    // refresh token (doing so would invalidate the lease holder's token →
+    // invalid_grant → bricked account). It serves on its existing access token
+    // for the bounded drain; the lease holder owns all rotation.
+    if (!this.writerLease) return true;
+
     if (!force && !isTokenExpiringSoon(account.expiresAt)) return true;
 
     // Coalesce concurrent refreshes
@@ -1642,6 +1745,9 @@ export class AccountManager {
 
     account._refreshPromise = (async () => {
       console.log(`[Maxpool] Refreshing token for account "${account.name}"...`);
+      // Record the token we're rotating FROM so the persistence layer's
+      // generation guard can detect another writer having already advanced it.
+      account._refreshedFrom = account.refreshToken;
       try {
         const newTokens = await this._refreshAccessToken(account.refreshToken);
         account.credential = newTokens.accessToken;
@@ -1673,6 +1779,18 @@ export class AccountManager {
     return account._refreshPromise;
   }
 
+  /**
+   * Await every in-flight OAuth token refresh to settle. The single-writer baton
+   * uses this on RELEASE: a refresh that passed the `if(!writerLease) return` gate
+   * BEFORE the lease was dropped is still awaiting its OAuth POST; the new worker
+   * must not acquire the lease and rotate the SAME single-use token until these
+   * settle, or the upstream invalidates one token → invalid_grant → bricked.
+   */
+  async drainRefreshes() {
+    const pending = this.accounts.map(a => a._refreshPromise).filter(Boolean);
+    if (pending.length) await Promise.allSettled(pending);
+  }
+
   /**
    * Set a callback to persist refreshed tokens to config.
    */

package/src/config.js

@@ -171,28 +171,100 @@ export async function loadState() {
   }
 }
 
-export async function saveState(state) {
-  await atomicWrite(getStatePath(), JSON.stringify(state, null, 2) + '\n');
+/**
+ * Read just the `_generation` integer of an on-disk JSON file (config or state)
+ * without parsing it as a typed object. Returns 0 when the file is missing,
+ * unreadable, or carries no generation (legacy files), so a first write always
+ * advances to >=1. Never throws.
+ */
+export async function readGeneration(path) {
+  try {
+    const parsed = JSON.parse(await readFile(path, 'utf-8'));
+    const g = Number(parsed?._generation);
+    return Number.isFinite(g) && g > 0 ? g : 0;
+  } catch {
+    return 0;
+  }
+}
+
+/**
+ * Generation-guarded state write (defense-in-depth single-writer enforcement).
+ *
+ * `state.quota` carries the learned quota. We stamp a monotonic `_generation`
+ * and refuse to clobber a file whose on-disk generation is NEWER than the one
+ * we last observed — that means another writer (a stale ex-lease worker) raced
+ * us, and overwriting would revert fresher quota. The baton sequencing already
+ * guarantees a single writer; this guard catches a sequencing bug.
+ *
+ * Pass `{ expectedGeneration }` to assert the on-disk file is still at the
+ * generation you read; omit it to read-then-bump unconditionally (single-owner
+ * fast path). Returns the generation written, or null when refused.
+ */
+let _stateWriteChain = Promise.resolve();
+
+export function saveState(state, { expectedGeneration = null } = {}) {
+  // Serialize state writes (a 60s interval flush can otherwise race the final
+  // baton flush, read the same on-disk generation, and double-write).
+  const run = async () => {
+    const path = getStatePath();
+    const onDisk = await readGeneration(path);
+    if (expectedGeneration != null && onDisk > expectedGeneration) {
+      // A newer writer advanced the file under us — refuse the stale flush.
+      return null;
+    }
+    const next = onDisk + 1;
+    await atomicWrite(path, JSON.stringify({ ...state, _generation: next }, null, 2) + '\n');
+    return next;
+  };
+  const result = _stateWriteChain.then(run, run);
+  _stateWriteChain = result.then(() => {}, () => {});
+  return result;
 }
 
+/** Barrier: resolves once every queued state write has settled. */
+export const flushStateWrites = () => _stateWriteChain;
+
 let _configWriteChain = Promise.resolve();
 
+/** Barrier: resolves once every queued config write (e.g. a fire-and-forget
+ *  token-refresh persist) has settled. Awaited on baton release so a rotated
+ *  token can't be left unwritten when the new worker boots from disk (M3). */
+export const flushConfigWrites = () => _configWriteChain;
+
 /**
  * Serialize an in-process read-modify-write of the config so concurrent updates
  * cannot lose each other's writes — e.g. a background OAuth token refresh
  * (fire-and-forget) racing a TUI account add/delete. Each update re-reads the
  * latest config, applies updater(config), then saves atomically.
  *
+ * Defense-in-depth single-writer guard: the config carries a monotonic
+ * `_generation`. If `guardGeneration` is supplied and the on-disk generation is
+ * already NEWER than it, the write is REFUSED (a stale ex-lease worker raced the
+ * current writer). The updater additionally receives the on-disk generation as
+ * `config._generation` so a refresh updater can SKIP a token rotation it sees a
+ * fresher writer already performed. The generation is bumped on every write.
+ *
  * NOTE: this serializes writes within THIS process only. A separate
  * `maxpool import`/`login` process writing concurrently is not coordinated
  * (that would require a lockfile) — but those are short, rare, human-driven.
  */
-export function atomicConfigUpdate(updater) {
+export function atomicConfigUpdate(updater, { guardGeneration = null } = {}) {
   const run = async () => {
     const config = await loadConfig() || createDefaultConfig();
-    await updater(config);
+    const onDisk = Number(config._generation) || 0;
+    if (guardGeneration != null && onDisk > guardGeneration) {
+      // Refuse: a newer writer already advanced the config. Surface the live
+      // config so the caller can reconcile, but write nothing.
+      const refused = new Error('config generation advanced; stale write refused');
+      refused.code = 'STALE_GENERATION';
+      refused.onDiskGeneration = onDisk;
+      refused.config = config;
+      throw refused;
+    }
+    const result = await updater(config);
+    config._generation = onDisk + 1;
     await saveConfig(config);
-    return config;
+    return result === undefined ? config : result;
   };
   // Run after any in-flight update regardless of whether it resolved or
   // rejected, so one failure doesn't poison the chain; the caller still sees