n9router 0.3.91 → 0.3.94

package/.next/standalone/src/mitm/server.js

@@ -7,7 +7,7 @@ const { log, err } = require("./logger");
 const { TARGET_HOSTS, URL_PATTERNS, getToolForHost } = require("./config");
 const { DATA_DIR, MITM_DIR } = require("./paths");
 const { isTokenSwapEnabled, getAllActiveConnections, triggerRefreshIfNeeded,
-        setCooldown, setModelCooldown, getTokenSwapStrategy,
+        forceRefreshConnection, setCooldown, setAuthCooldown, setModelCooldown, getTokenSwapStrategy,
         parseQuotaCooldown, markAccountUsed, getConnectionLabel } = require("./tokenPool");
 const { getCertForDomain } = require("./cert/generate");
 const { buildInputOnlyRequestDetail, createTokenSwapUsageObserver, generateDetailId } = require("./usageTracker");
@@ -187,120 +187,144 @@ async function tokenSwapForward(req, res, bodyBuffer, connections, model, strate
 
   for (let i = 0; i < connections.length; i++) {
     const originalConn = connections[i];
-    const conn = await triggerRefreshIfNeeded(originalConn);
-
-    const label = getConnectionLabel(conn);
-    const modelTag = model ? ` model=${model}` : "";
-    const posTag = connections.length > 1 ? ` [${i + 1}/${connections.length}]` : "";
-    const useTag = conn.consecutiveUseCount > 1 ? ` uses=${conn.consecutiveUseCount}` : "";
-    log(`🔑 [token-swap]${posTag} trying "${label}"${modelTag}${useTag}`);
-
-    const swappedHeaders = {
-      ...req.headers,
-      host: targetHost,
-      authorization: `Bearer ${conn.accessToken}`
-    };
-
-    try {
-      const result = await new Promise((resolve, reject) => {
-        const forwardReq = https.request({
-          hostname: targetIP,
-          port: 443,
-          path: req.url,
-          method: req.method,
-          headers: swappedHeaders,
-          servername: targetHost,
-          rejectUnauthorized: false
-        }, (forwardRes) => {
-          if (forwardRes.statusCode === 429 || forwardRes.statusCode === 503) {
-            // Buffer the short error body to parse cooldown duration
-            const chunks = [];
-            forwardRes.on("data", c => chunks.push(c));
-            forwardRes.on("end", () => {
-              const body = Buffer.concat(chunks).toString();
-              resolve({ retry: true, body, statusCode: forwardRes.statusCode });
-            });
+    let conn = await triggerRefreshIfNeeded(originalConn);
+    let authRefreshAttempted = false;
+
+    while (true) {
+      const label = getConnectionLabel(conn);
+      const modelTag = model ? ` model=${model}` : "";
+      const posTag = connections.length > 1 ? ` [${i + 1}/${connections.length}]` : "";
+      const useTag = conn.consecutiveUseCount > 1 ? ` uses=${conn.consecutiveUseCount}` : "";
+      log(`🔑 [token-swap]${posTag} trying "${label}"${modelTag}${useTag}`);
+
+      const swappedHeaders = {
+        ...req.headers,
+        host: targetHost,
+        authorization: `Bearer ${conn.accessToken}`
+      };
+
+      try {
+        const result = await new Promise((resolve, reject) => {
+          const forwardReq = https.request({
+            hostname: targetIP,
+            port: 443,
+            path: req.url,
+            method: req.method,
+            headers: swappedHeaders,
+            servername: targetHost,
+            rejectUnauthorized: false
+          }, (forwardRes) => {
+            if (forwardRes.statusCode === 429 || forwardRes.statusCode === 503 || forwardRes.statusCode === 401) {
+              const chunks = [];
+              forwardRes.on("data", c => chunks.push(c));
+              forwardRes.on("end", () => {
+                const body = Buffer.concat(chunks).toString();
+                const retryType = forwardRes.statusCode === 401 ? "auth" : "quota";
+                resolve({ retry: true, retryType, body, headers: forwardRes.headers, statusCode: forwardRes.statusCode });
+              });
+            } else {
+              resolve({ retry: false, response: forwardRes });
+            }
+          });
+          forwardReq.on("error", reject);
+          if (bodyBuffer.length > 0) forwardReq.write(bodyBuffer);
+          forwardReq.end();
+        });
+
+        if (result.retry && result.retryType === "quota") {
+          const cooldownMs = parseQuotaCooldown(result.body);
+          const cdLabel = cooldownMs ? ` cooldown=${Math.ceil(cooldownMs / 60000)}m` : "";
+          if (strategy === "sticky" && model) {
+            setModelCooldown(conn.id, model, cooldownMs);
+            log(`⚠️ [token-swap] "${label}" → ${result.statusCode} model=${model} quota exhausted${cdLabel}, trying next...`);
           } else {
-            // Success or non-quota error → pipe to client
-            resolve({ retry: false, response: forwardRes });
+            setCooldown(conn.id, cooldownMs);
+            log(`⚠️ [token-swap] "${label}" → ${result.statusCode} account quota exhausted${cdLabel}, trying next...`);
           }
-        });
-        forwardReq.on("error", reject);
-        if (bodyBuffer.length > 0) forwardReq.write(bodyBuffer);
-        forwardReq.end();
-      });
-
-      if (result.retry) {
-        const cooldownMs = parseQuotaCooldown(result.body);
-        const cdLabel = cooldownMs ? ` cooldown=${Math.ceil(cooldownMs / 60000)}m` : "";
-        if (strategy === "sticky" && model) {
-          // Sticky: mark only this account+model as exhausted, not the whole account
-          setModelCooldown(conn.id, model, cooldownMs);
-          log(`⚠️ [token-swap] "${label}" → ${result.statusCode} model=${model} quota exhausted${cdLabel}, trying next...`);
-        } else {
-          // Round-robin: put whole account on cooldown
-          setCooldown(conn.id, cooldownMs);
-          log(`⚠️ [token-swap] "${label}" → ${result.statusCode} account quota exhausted${cdLabel}, trying next...`);
+          break;
         }
-        continue;
-      }
 
-      const newCount = (conn.consecutiveUseCount || 0) + 1;
-      const statusCode = result.response.statusCode || 0;
-      const successModelTag = model ? ` model=${model}` : "";
-      const successStrategyTag = strategy === "sticky" ? ` sticky(use #${newCount})` : ` rr`;
-      log(`✅ [token-swap] "${label}" → ${statusCode}${successModelTag}${successStrategyTag}`);
-      markAccountUsed(conn.id);
-      res.writeHead(statusCode, result.response.headers);
-
-      const detailId = generateDetailId(model);
-      const inputOnlyDetail = buildInputOnlyRequestDetail({
-        detailId,
-        provider,
-        model,
-        connectionId: conn.id,
-        bodyBuffer
-      });
-      fetch("http://127.0.0.1:20128/api/internal/request-detail", {
-        method: "POST",
-        headers: {
-          "Content-Type": "application/json",
-          [INTERNAL_REQUEST_HEADER.name]: INTERNAL_REQUEST_HEADER.value
-        },
-        body: JSON.stringify(inputOnlyDetail)
-      }).catch((detailError) => {
-        err(`[token-swap] failed to create request detail for "${label}": ${detailError.message}`);
-      });
-
-      const usageObserver = createTokenSwapUsageObserver({
-        provider,
-        model,
-        connectionId: conn.id,
-        accountLabel: label,
-        bodyBuffer,
-        contentType: result.response.headers["content-type"] || "",
-        contentEncoding: result.response.headers["content-encoding"] || "",
-        statusCode,
-        detailRecord: inputOnlyDetail,
-        requestStartTime
-      });
-
-      result.response.on("data", (chunk) => {
-        usageObserver.onChunk(chunk);
-        res.write(chunk);
-      });
-      result.response.on("end", () => {
-        res.end();
-        usageObserver.onEnd().catch(() => {});
-      });
-      result.response.on("error", (streamError) => {
-        err(`[token-swap] upstream stream error for "${label}": ${streamError.message}`);
-        if (!res.writableEnded) res.end();
-      });
-      return true;
-    } catch (e) {
-      err(`[token-swap] error for "${label}": ${e.message}`);
-      continue;
+        if (result.retry && result.retryType === "auth") {
+          const retryableAuth = isRetryableAuthFailure(result.statusCode, result.headers, result.body);
+          if (!retryableAuth) {
+            res.writeHead(result.statusCode, result.headers);
+            res.end(result.body);
+            return true;
+          }
+
+          if (!authRefreshAttempted && conn.refreshToken) {
+            authRefreshAttempted = true;
+            log(`⚠️ [token-swap] "${label}" → 401 invalid_token, forcing refresh...`);
+            const refreshResult = await forceRefreshConnection(conn);
+            conn = refreshResult.connection || conn;
+            if (refreshResult.refreshed) {
+              log(`↻ [token-swap] "${label}" refreshed, retrying same account...`);
+              continue;
+            }
+          }
+
+          setAuthCooldown(conn.id);
+          log(`⚠️ [token-swap] "${label}" → 401 invalid_token, trying next...`);
+          break;
+        }
+
+        const newCount = (conn.consecutiveUseCount || 0) + 1;
+        const statusCode = result.response.statusCode || 0;
+        const successModelTag = model ? ` model=${model}` : "";
+        const successStrategyTag = strategy === "sticky" ? ` sticky(use #${newCount})` : ` rr`;
+        log(`✅ [token-swap] "${label}" → ${statusCode}${successModelTag}${successStrategyTag}`);
+        markAccountUsed(conn.id);
+        res.writeHead(statusCode, result.response.headers);
+
+        const detailId = generateDetailId(model);
+        const inputOnlyDetail = buildInputOnlyRequestDetail({
+          detailId,
+          provider,
+          model,
+          connectionId: conn.id,
+          bodyBuffer
+        });
+        fetch("http://127.0.0.1:20128/api/internal/request-detail", {
+          method: "POST",
+          headers: {
+            "Content-Type": "application/json",
+            [INTERNAL_REQUEST_HEADER.name]: INTERNAL_REQUEST_HEADER.value
+          },
+          body: JSON.stringify(inputOnlyDetail)
+        }).catch((detailError) => {
+          err(`[token-swap] failed to create request detail for "${label}": ${detailError.message}`);
+        });
+
+        const usageObserver = createTokenSwapUsageObserver({
+          provider,
+          model,
+          connectionId: conn.id,
+          accountLabel: label,
+          bodyBuffer,
+          contentType: result.response.headers["content-type"] || "",
+          contentEncoding: result.response.headers["content-encoding"] || "",
+          statusCode,
+          detailRecord: inputOnlyDetail,
+          requestStartTime
+        });
+
+        result.response.on("data", (chunk) => {
+          usageObserver.onChunk(chunk);
+          res.write(chunk);
+        });
+        result.response.on("end", () => {
+          res.end();
+          usageObserver.onEnd().catch(() => {});
+        });
+        result.response.on("error", (streamError) => {
+          err(`[token-swap] upstream stream error for "${label}": ${streamError.message}`);
+          if (!res.writableEnded) res.end();
+        });
+        return true;
+      } catch (e) {
+        err(`[token-swap] error for "${label}": ${e.message}`);
+        break;
+      }
     }
   }
 
@@ -308,6 +332,36 @@ async function tokenSwapForward(req, res, bodyBuffer, connections, model, strate
   return false;
 }
 
+function getHeaderValue(headers, name) {
+  const value = headers?.[String(name).toLowerCase()];
+  if (Array.isArray(value)) return value.join(", ");
+  return typeof value === "string" ? value : "";
+}
+
+function isRetryableAuthFailure(statusCode, headers, body) {
+  if (statusCode !== 401) return false;
+
+  const authHeader = getHeaderValue(headers, "www-authenticate").toLowerCase();
+  if (authHeader.includes("invalid_token")) return true;
+
+  try {
+    const parsed = JSON.parse(body || "{}");
+    const status = String(parsed?.error?.status || parsed?.status || "").toUpperCase();
+    const message = String(parsed?.error?.message || parsed?.message || "").toLowerCase();
+    if (status === "UNAUTHENTICATED") return true;
+    if (message.includes("invalid authentication credentials")) return true;
+    if (message.includes("invalid token")) return true;
+    if (message.includes("unauthenticated")) return true;
+  } catch {
+    const fallback = String(body || "").toLowerCase();
+    if (fallback.includes("invalid authentication credentials")) return true;
+    if (fallback.includes("invalid token")) return true;
+    if (fallback.includes("unauthenticated")) return true;
+  }
+
+  return false;
+}
+
 // ── Request handler ───────────────────────────────────────────
 
 const server = https.createServer(sslOptions, async (req, res) => {

package/.next/standalone/src/mitm/tokenPool.js

@@ -14,6 +14,7 @@ const DB_FILE = path.join(DATA_DIR, "db.json");
 const TOKEN_EXPIRY_BUFFER_MS = 5 * 60 * 1000; // refresh 5min before expiry
 const ROUTER_PORT = process.env.PORT || 20128;
 const DEFAULT_COOLDOWN_MS = 5 * 60 * 1000;
+const DEFAULT_AUTH_COOLDOWN_MS = 10 * 60 * 1000;
 
 // ── In-memory state ──────────────────────────────────────────
 const cooldownMap = {};        // { [connectionId]: expiresTimestamp }
@@ -28,6 +29,12 @@ function setCooldown(connId, durationMs) {
   log(`⏸ [token-pool] cooldown: ${connId.slice(0, 8)}… for ${Math.ceil(ms / 60000)}m`);
 }
 
+function setAuthCooldown(connId, durationMs) {
+  const ms = durationMs || DEFAULT_AUTH_COOLDOWN_MS;
+  cooldownMap[connId] = Date.now() + ms;
+  log(`🔒 [token-pool] auth-cooldown: ${connId.slice(0, 8)}… for ${Math.ceil(ms / 60000)}m`);
+}
+
 function isInCooldown(connId) {
   if (!cooldownMap[connId]) return false;
   if (Date.now() > cooldownMap[connId]) {
@@ -287,41 +294,82 @@ function markAccountUsed(connId) {
 // checks expiry, refreshes token automatically, then returns the
 // latest persisted connection snapshot for the current request.
 
-async function triggerRefreshIfNeeded(connection) {
-  if (!connection?.expiresAt || !connection?.refreshToken) return connection;
-  const expiresAt = new Date(connection.expiresAt).getTime();
-  if (Date.now() + TOKEN_EXPIRY_BUFFER_MS < expiresAt) return connection;
-
-  log(`🔄 [token-pool] near-expiry refresh → ${getConnectionLabel(connection).slice(0, 20)}`);
+async function runConnectionTest(connection) {
   return await new Promise((resolve) => {
+    if (!connection?.id) {
+      resolve({ connection, refreshed: false, valid: false });
+      return;
+    }
+
+    let responseBody = "";
+
     try {
       const req = http.request(
         { hostname: "127.0.0.1", port: ROUTER_PORT, path: `/api/providers/${connection.id}/test`, method: "POST" },
         (res) => {
-          res.on("data", () => {});
+          res.on("data", (chunk) => {
+            responseBody += chunk.toString();
+          });
           res.on("end", () => {
-            const refreshed = getConnectionById(connection.id);
-            if (refreshed?.accessToken && refreshed.accessToken !== connection.accessToken) {
-              log(`♻️ [token-pool] refreshed token applied → ${getConnectionLabel(connection).slice(0, 20)}`);
+            let payload = {};
+            try {
+              payload = JSON.parse(responseBody || "{}");
+            } catch {
+              payload = {};
             }
-            resolve(refreshed || connection);
+
+            const refreshedConnection = getConnectionById(connection.id);
+            const tokenChanged = !!(refreshedConnection?.accessToken && refreshedConnection.accessToken !== connection.accessToken);
+            resolve({
+              connection: refreshedConnection || connection,
+              refreshed: !!payload.refreshed || tokenChanged,
+              valid: !!payload.valid,
+            });
           });
         }
       );
-      req.on("error", () => resolve(connection));
+      req.on("error", () => resolve({ connection, refreshed: false, valid: false }));
       req.end();
     } catch {
-      resolve(connection);
+      resolve({ connection, refreshed: false, valid: false });
     }
   });
 }
 
+async function triggerRefreshIfNeeded(connection) {
+  if (!connection?.expiresAt || !connection?.refreshToken) return connection;
+  const expiresAt = new Date(connection.expiresAt).getTime();
+  if (Date.now() + TOKEN_EXPIRY_BUFFER_MS < expiresAt) return connection;
+
+  log(`🔄 [token-pool] near-expiry refresh → ${getConnectionLabel(connection).slice(0, 20)}`);
+  const result = await runConnectionTest(connection);
+  if (result.refreshed) {
+    log(`♻️ [token-pool] refreshed token applied → ${getConnectionLabel(connection).slice(0, 20)}`);
+  }
+  return result.connection || connection;
+}
+
+async function forceRefreshConnection(connection) {
+  if (!connection?.refreshToken) {
+    return { connection, refreshed: false, valid: false };
+  }
+
+  log(`🔄 [token-pool] auth refresh → ${getConnectionLabel(connection).slice(0, 20)}`);
+  const result = await runConnectionTest(connection);
+  if (result.refreshed) {
+    log(`♻️ [token-pool] refreshed token applied → ${getConnectionLabel(connection).slice(0, 20)}`);
+  }
+  return result;
+}
+
 module.exports = {
   isTokenSwapEnabled,
   getNextConnection,
   getAllActiveConnections,
   triggerRefreshIfNeeded,
+  forceRefreshConnection,
   setCooldown,
+  setAuthCooldown,
   setModelCooldown,
   isModelExhausted,
   getTokenSwapStrategy,