maxpool 1.3.1 → 1.5.0

package/src/oauth.js

@@ -4,7 +4,10 @@ import { createInterface } from 'node:readline';
 import http from 'node:http';
 
 const PROFILE_URL = 'https://api.anthropic.com/api/oauth/profile';
-const DEFAULT_TOKEN_ENDPOINT = 'https://platform.claude.com/v1/oauth/token';
+// Token endpoint is overridable via env so a stub OAuth server can be pointed at
+// for integration tests (e.g. the single-use-rotating-token reload torture test).
+const DEFAULT_TOKEN_ENDPOINT = process.env.MAXPOOL_OAUTH_TOKEN_ENDPOINT
+  || 'https://platform.claude.com/v1/oauth/token';
 const DEFAULT_CLIENT_ID = '9d1c250a-e61b-44d9-88ed-5944d1962f5e';
 
 /**

package/src/prober.js

@@ -40,20 +40,29 @@ export class Prober {
     }
   }
 
-  stop() {
+  /** Stop scheduling AND await any probe cycle already in flight, so a caller
+   *  (the baton release) can be sure no probe-driven token rotation is pending
+   *  before it hands the writer lease to another worker. */
+  async stop() {
     if (this.timer) { clearInterval(this.timer); this.timer = null; }
+    if (this._inflight) { try { await this._inflight; } catch { /* swallow */ } }
   }
 
-  /** Probe every OAuth account once. Overlapping cycles are skipped. */
-  async probeAll() {
-    if (this._running) return;
+  /** Probe every OAuth account once. Overlapping cycles are skipped. The active
+   *  cycle is tracked on `_inflight` so stop() can await it. */
+  probeAll() {
+    if (this._running) return this._inflight || Promise.resolve();
     this._running = true;
-    try {
-      const accounts = this.am.accounts.filter(a => a.type === 'oauth' && a.credential);
-      await Promise.all(accounts.map(a => this.probeOne(a)));
-    } finally {
-      this._running = false;
-    }
+    this._inflight = (async () => {
+      try {
+        const accounts = this.am.accounts.filter(a => a.type === 'oauth' && a.credential);
+        await Promise.all(accounts.map(a => this.probeOne(a)));
+      } finally {
+        this._running = false;
+        this._inflight = null;
+      }
+    })();
+    return this._inflight;
   }
 
   async probeOne(account) {

package/src/reload-protocol.js

@@ -0,0 +1,121 @@
+// IPC message vocabulary + the supervisor-side baton orchestration for the
+// near-zero-downtime reload (issue #6). Pure logic, no process side-effects, so
+// the baton sequence is unit-testable against fake workers.
+//
+// Roles:
+//   Supervisor — owns the listening socket for life, never closes it, spawns
+//     workers and hands the socket HANDLE to exactly one acceptor at a time.
+//   Worker     — accepts on the shared handle; exactly one worker holds the
+//     WRITER LEASE (refresh/probe/persist) at any instant.
+//
+// The baton (single-writer guarantee): old worker RELEASES (stops accepting +
+// stops writing, flushes once) BEFORE the new worker ACQUIRES. Reads overlap;
+// writes never do.
+
+// Supervisor → worker
+export const MSG_LISTEN = 'listen';       // (with handle) accept on this socket + take TUI/lease
+export const MSG_RELEASE = 'release';     // stop accepting, stop writing, flush, give up TUI
+export const MSG_TAKEOVER = 'takeover';   // (with handle) start accepting + acquire writer lease + TUI
+export const MSG_PROBE_READY = 'probe-ready'; // ask a headless worker to confirm it booted OK
+
+// Worker → supervisor
+export const MSG_RELOAD_REQUEST = 'reload-request'; // primary worker asks for a seamless reload
+export const MSG_READY = 'ready';         // headless worker booted new code successfully
+export const MSG_FAILED = 'failed';       // worker failed to boot/bind/restore
+export const MSG_RELEASED = 'released';   // old worker stopped accepting + writing, flushed
+export const MSG_PRIMARY = 'primary';     // new worker is now sole acceptor + lease holder
+export const MSG_QUOTA_STATE = 'quota-state'; // old worker hands its in-memory quota to supervisor
+
+/**
+ * Outcome codes for a reload attempt, surfaced to the supervisor so it can log
+ * and decide fallback. SWAPPED = fully cut over; ROLLED_BACK = old worker stays
+ * primary (no disruption); FALLBACK = abrupt exit-75 restart path taken.
+ */
+export const RELOAD_SWAPPED = 'swapped';
+export const RELOAD_ROLLED_BACK = 'rolled-back';
+export const RELOAD_FALLBACK = 'fallback';
+
+/**
+ * Drive the baton handshake over two worker "channels" (anything with
+ * .send(msg[,handle]) and an async waitFor(type, timeoutMs) that resolves with
+ * the message payload or rejects on timeout/exit). Returns one of the RELOAD_*
+ * outcomes. The caller (supervisor) owns spawning, killing, reaping and the
+ * actual fallback restart; this function only sequences the protocol and tells
+ * the caller what happened.
+ *
+ * Sequence (matches design step 2a–2e):
+ *   a. new worker already spawned headless (no lease); we ask it to confirm READY
+ *   b. READY fails  → ROLLED_BACK: caller kills new, old stays primary
+ *   c. old RELEASE  → waits RELEASED (old stopped accepting + writing, flushed)
+ *   d. new TAKEOVER → waits PRIMARY (new is sole acceptor + lease holder)
+ *   e. caller drains+reaps old
+ *
+ * Any throw → FALLBACK (caller does the tested abrupt restart). All-or-nothing:
+ * we never leave both accepting or both writing.
+ */
+export async function runReloadBaton({
+  oldWorker,
+  newWorker,
+  handle,
+  prepareHandle = async h => h,
+  readyTimeoutMs = 10_000,
+  releaseTimeoutMs = 20_000,
+  takeoverTimeoutMs = 10_000,
+  log = () => {},
+}) {
+  // (a) Confirm the freshly-spawned headless worker booted the new code.
+  let ready;
+  try {
+    newWorker.send({ type: MSG_PROBE_READY });
+    ready = await newWorker.waitFor([MSG_READY, MSG_FAILED], readyTimeoutMs);
+  } catch (err) {
+    log(`reload: readiness wait failed (${err.message}); rolling back`);
+    return RELOAD_ROLLED_BACK;
+  }
+  // (b) New worker did not come up cleanly → roll back fully, old stays primary.
+  if (!ready || ready.type === MSG_FAILED) {
+    log(`reload: new worker reported failed (${ready?.reason || 'no ready'}); rolling back`);
+    return RELOAD_ROLLED_BACK;
+  }
+
+  // (c) Old worker releases: stop accepting NEW, keep in-flight, stop writing,
+  //     flush config+state ONE final time, drop the TUI. After this point the
+  //     old worker holds no lease and no acceptor — and (single-writer) no
+  //     refresh/prober write is still in flight (it awaits drainRefreshes).
+  try {
+    oldWorker.send({ type: MSG_RELEASE });
+    await oldWorker.waitFor([MSG_RELEASED], releaseTimeoutMs);
+  } catch (err) {
+    // Old worker didn't ack release. We must NOT hand the socket to the new
+    // worker (would risk two acceptors) — fall back to the abrupt path.
+    log(`reload: old worker did not release (${err.message}); falling back`);
+    return RELOAD_FALLBACK;
+  }
+
+  // The old worker has now freed the listening fd. Re-arm the supervisor's own
+  // acceptor over the cutover gap (covers new conns in the OS backlog) and get a
+  // fresh re-sendable handle for the new worker.
+  let liveHandle = handle;
+  try {
+    liveHandle = await prepareHandle(handle);
+  } catch (err) {
+    log(`reload: could not re-arm listening socket (${err.message}); falling back`);
+    return RELOAD_FALLBACK;
+  }
+
+  // (d) New worker takes over: it becomes the SOLE acceptor (old already closed)
+  //     and ACQUIRES the writer lease (refresh/probe/persist re-enabled) + TUI.
+  try {
+    newWorker.send({ type: MSG_TAKEOVER }, liveHandle);
+    await newWorker.waitFor([MSG_PRIMARY], takeoverTimeoutMs);
+  } catch (err) {
+    // The new worker failed to start accepting after the old already released.
+    // Neither is accepting now — fall back so the supervisor force-restarts and
+    // a worker comes back up on the supervisor-owned socket (queued conns drain).
+    log(`reload: new worker did not take over (${err.message}); falling back`);
+    return RELOAD_FALLBACK;
+  }
+
+  log('reload: cutover complete; new worker is primary');
+  return RELOAD_SWAPPED;
+}

package/src/server.js

@@ -51,8 +51,25 @@ export function createProxyServer(accountManager, config, hooks = {}) {
     mkdir(logDir, { recursive: true }).catch(() => {});
   }
 
+  // Reload drain flag. When the worker releases the baton it sets this so every
+  // remaining response carries `Connection: close`, retiring the client's
+  // keep-alive socket instead of letting it pipeline a NEW request onto a worker
+  // that's shutting down. `connection` is hop-by-hop so it's stripped from the
+  // upstream response headers — a setHeader here survives the later writeHead.
+  let draining = false;
+
+  // Identifies the WORKER process that served a response — proves the supervisor
+  // (which holds the socket but does not serve) never swallowed the request. Set
+  // before any writeHead; `x-maxpool-*` is informative-only and stripped from
+  // upstream-bound request headers elsewhere.
+  const workerStamp = String(process.pid);
+
   const server = http.createServer(async (req, res) => {
     try {
+      try { res.setHeader('x-maxpool-worker', workerStamp); } catch { /* headers sent */ }
+      if (draining) {
+        try { res.setHeader('Connection', 'close'); } catch { /* headers may be sent */ }
+      }
       // Auth check — skip for localhost connections
       const clientKey = req.headers['x-api-key'];
       const remoteAddr = req.socket.remoteAddress;
@@ -187,6 +204,11 @@ export function createProxyServer(accountManager, config, hooks = {}) {
     }
   });
 
+  // Begin reload drain: every subsequent response gets `Connection: close` so
+  // keep-alive clients retire their socket and don't pipeline a new request onto
+  // this releasing worker.
+  server.maxpoolBeginDrain = () => { draining = true; };
+
   return server;
 }
 
@@ -930,7 +952,7 @@ function unavailableMessage(accountManager, requestInfo = {}, retryAfter, willRe
   return `All ${n} accounts exhausted. Retry in ${retryAfter}s.`;
 }
 
-export const __serverTest = { unavailableMessage, isRetriableUpstreamStatus, headerValue, getMaxpoolProfile, ensureQueueHeartbeat, clearQueueHeartbeat };
+export const __serverTest = { unavailableMessage, isRetriableUpstreamStatus, headerValue, getMaxpoolProfile, ensureQueueHeartbeat, clearQueueHeartbeat, describeRequest };
 
 async function readErrorBody(upstreamRes, limitBytes = 64 * 1024) {
   if (!upstreamRes.body) return '';
@@ -1375,6 +1397,11 @@ function describeRequest(req, body) {
     if (requiresAnthropicThinkingIntegrity(json)) {
       info.requiresAnthropicThinkingIntegrity = true;
     }
+    // We fully scanned this body for signed-thinking content. Only a successfully
+    // scanned, thinking-free body is safe to migrate to another account (session
+    // rebalancing); an unparsed body leaves this false → treated as NOT safe
+    // (fail-closed) so we never replay a signed thinking block to a new account.
+    info.bodyThinkingScanned = true;
   } catch {
     // Non-JSON requests are rare; body size still gives a useful load signal.
   }