feat: add refresh retry/timeout and tests

2025-12-27 15:07:54 +08:00
parent 5e01993120
commit a18e45ee78
6 changed files with 331 additions and 85 deletions
--- a/.env.example
+++ b/.env.example
@@ -6,6 +6,11 @@ FACTORY_API_KEY=your_factory_api_key_here
 # 方式2：使用refresh token自动刷新（次优先级）
 DROID_REFRESH_KEY=your_refresh_token_here

+# refresh token 请求超时与重试（可选）
+DROID_REFRESH_TIMEOUT_MS=15000
+DROID_REFRESH_RETRIES=2
+DROID_REFRESH_RETRY_BASE_MS=500
+
 # 阿里云日志服务配置
 SLS_ENABLED=false
 ALIYUN_ACCESS_KEY_ID=your_access_key_id
--- a/account-manager.js
+++ b/account-manager.js
@@ -1,8 +1,8 @@
 import fs from 'fs';
 import path from 'path';
-import fetch from 'node-fetch';
 import { logInfo, logDebug, logError } from './logger.js';
 import { getNextProxyAgent } from './proxy-manager.js';
+import { getRefreshConfig, requestRefreshToken } from './refresh-client.js';

 /**
 * Account Manager - 管理多个 OAuth 账号的选择、刷新和统计
@@ -21,6 +21,7 @@ const CLIENT_ID = 'client_01HNM792M5G5G1A2THWPXKFMXB';
 class AccountManager {
    constructor() {
        this.accounts = [];
+        this.refreshLocks = new Map(); // 刷新锁，避免同一账号并发刷新
        this.settings = {
            algorithm: 'weighted', // 'weighted' or 'simple'
            refresh_interval_hours: REFRESH_INTERVAL_HOURS,
@@ -174,12 +175,30 @@ class AccountManager {
        const needsRefresh = !account.access_token || this._shouldRefresh(account);

        if (needsRefresh) {
-            await this._refreshToken(account);
+            await this._refreshTokenWithLock(account);
        }

        return account.access_token;
    }

+    /**
+     * 带锁刷新，避免同一账号并发刷新
+     */
+    async _refreshTokenWithLock(account) {
+        const existing = this.refreshLocks.get(account.id);
+        if (existing) {
+            return existing;
+        }
+
+        const refreshPromise = this._refreshToken(account)
+            .finally(() => {
+                this.refreshLocks.delete(account.id);
+            });
+
+        this.refreshLocks.set(account.id, refreshPromise);
+        return refreshPromise;
+    }
+
    /**
     * 检查是否需要刷新 token
     */
@@ -198,32 +217,15 @@ class AccountManager {
        logInfo(`AccountManager: 刷新账号 ${account.id} 的 token...`);

        try {
-            const formData = new URLSearchParams();
-            formData.append('grant_type', 'refresh_token');
-            formData.append('refresh_token', account.refresh_token);
-            formData.append('client_id', CLIENT_ID);
-
            const proxyAgentInfo = getNextProxyAgent(REFRESH_URL);
-            const fetchOptions = {
-                method: 'POST',
-                headers: {
-                    'Content-Type': 'application/x-www-form-urlencoded'
-                },
-                body: formData.toString()
-            };
-
-            if (proxyAgentInfo?.agent) {
-                fetchOptions.agent = proxyAgentInfo.agent;
-            }
-
-            const response = await fetch(REFRESH_URL, fetchOptions);
-
-            if (!response.ok) {
-                const errorText = await response.text();
-                throw new Error(`刷新失败: ${response.status} ${errorText}`);
-            }
-
-            const data = await response.json();
+            const refreshConfig = getRefreshConfig();
+            const data = await requestRefreshToken({
+                refreshUrl: REFRESH_URL,
+                refreshToken: account.refresh_token,
+                clientId: CLIENT_ID,
+                proxyAgentInfo,
+                ...refreshConfig
+            });

            // 更新账号信息
            account.access_token = data.access_token;
--- a/auth.js
+++ b/auth.js
@@ -1,10 +1,10 @@
 import fs from 'fs';
 import path from 'path';
 import os from 'os';
-import fetch from 'node-fetch';
 import { logDebug, logError, logInfo } from './logger.js';
 import { getNextProxyAgent } from './proxy-manager.js';
 import { getAccountManager, initializeAccountManager } from './account-manager.js';
+import { getRefreshConfig, requestRefreshToken } from './refresh-client.js';

 // State management for API key and refresh
 let currentApiKey = null;
@@ -16,6 +16,7 @@ let authFilePath = null;
 let factoryApiKey = null; // From FACTORY_API_KEY environment variable
 let multiAccountMode = false; // 是否启用多账号模式
 let lastSelectedAccountId = null; // 记录最后选择的账号ID，用于结果回调
+let refreshInFlight = null; // 刷新锁，避免并发刷新

 const REFRESH_URL = 'https://api.workos.com/user_management/authenticate';
 const REFRESH_INTERVAL_HOURS = 6; // Refresh every 6 hours
@@ -148,64 +149,58 @@ async function refreshApiKey() {
    throw new Error('No refresh token available');
  }

-  if (!clientId) {
-    clientId = 'client_01HNM792M5G5G1A2THWPXKFMXB';
-    logDebug(`Using fixed client ID: ${clientId}`);
+  if (refreshInFlight) {
+    return refreshInFlight;
  }

-  logInfo('Refreshing API key...');
+  refreshInFlight = (async () => {
+    if (!clientId) {
+      clientId = 'client_01HNM792M5G5G1A2THWPXKFMXB';
+      logDebug(`Using fixed client ID: ${clientId}`);
+    }
+
+    logInfo('Refreshing API key...');
+
+    try {
+    const proxyAgentInfo = getNextProxyAgent(REFRESH_URL);
+    const refreshConfig = getRefreshConfig();
+    const data = await requestRefreshToken({
+      refreshUrl: REFRESH_URL,
+      refreshToken: currentRefreshToken,
+      clientId,
+      proxyAgentInfo,
+      ...refreshConfig
+    });
+
+      // Update tokens
+      currentApiKey = data.access_token;
+      currentRefreshToken = data.refresh_token;
+      lastRefreshTime = Date.now();
+
+      // Log user info
+      if (data.user) {
+        logInfo(`Authenticated as: ${data.user.email} (${data.user.first_name} ${data.user.last_name})`);
+        logInfo(`User ID: ${data.user.id}`);
+        logInfo(`Organization ID: ${data.organization_id}`);
+      }
+
+      // Save tokens to file
+      saveTokens(data.access_token, data.refresh_token);
+
+      logInfo(`New Refresh-Key: ${currentRefreshToken}`);
+      logInfo('API key refreshed successfully');
+      return data.access_token;
+
+    } catch (error) {
+      logError('Failed to refresh API key', error);
+      throw error;
+    }
+  })();

  try {
-    // Create form data
-    const formData = new URLSearchParams();
-    formData.append('grant_type', 'refresh_token');
-    formData.append('refresh_token', currentRefreshToken);
-    formData.append('client_id', clientId);
-
-    const proxyAgentInfo = getNextProxyAgent(REFRESH_URL);
-    const fetchOptions = {
-      method: 'POST',
-      headers: {
-        'Content-Type': 'application/x-www-form-urlencoded'
-      },
-      body: formData.toString()
-    };
-
-    if (proxyAgentInfo?.agent) {
-      fetchOptions.agent = proxyAgentInfo.agent;
-    }
-
-    const response = await fetch(REFRESH_URL, fetchOptions);
-
-    if (!response.ok) {
-      const errorText = await response.text();
-      throw new Error(`Failed to refresh token: ${response.status} ${errorText}`);
-    }
-
-    const data = await response.json();
-
-    // Update tokens
-    currentApiKey = data.access_token;
-    currentRefreshToken = data.refresh_token;
-    lastRefreshTime = Date.now();
-
-    // Log user info
-    if (data.user) {
-      logInfo(`Authenticated as: ${data.user.email} (${data.user.first_name} ${data.user.last_name})`);
-      logInfo(`User ID: ${data.user.id}`);
-      logInfo(`Organization ID: ${data.organization_id}`);
-    }
-
-    // Save tokens to file
-    saveTokens(data.access_token, data.refresh_token);
-
-    logInfo(`New Refresh-Key: ${currentRefreshToken}`);
-    logInfo('API key refreshed successfully');
-    return data.access_token;
-
-  } catch (error) {
-    logError('Failed to refresh API key', error);
-    throw error;
+    return await refreshInFlight;
+  } finally {
+    refreshInFlight = null;
  }
 }

@@ -383,4 +378,3 @@ export function getAuthStatus() {
    multiAccountMode: false
  };
 }
-
--- a/package.json
+++ b/package.json
@@ -6,7 +6,8 @@
  "type": "module",
  "scripts": {
    "start": "node server.js",
-    "dev": "node server.js"
+    "dev": "node server.js",
+    "test": "node --test"
  },
  "keywords": [
    "openai",
@@ -21,4 +22,4 @@
    "node-fetch": "^3.3.2",
    "aliyun-log": "github:aliyun/aliyun-log-nodejs-sdk"
  }
-}
+}
--- a/refresh-client.js
+++ b/refresh-client.js
@@ -0,0 +1,132 @@
+import fetch from 'node-fetch';
+
+const DEFAULT_TIMEOUT_MS = 15000;
+const DEFAULT_MAX_RETRIES = 2;
+const DEFAULT_RETRY_BASE_MS = 500;
+const MAX_RETRY_DELAY_MS = 5000;
+
+function normalizeNumber(value, fallback) {
+  const parsed = parseInt(value, 10);
+  if (Number.isFinite(parsed) && parsed >= 0) {
+    return parsed;
+  }
+  return fallback;
+}
+
+export function getRefreshConfig() {
+  return {
+    timeoutMs: normalizeNumber(process.env.DROID_REFRESH_TIMEOUT_MS, DEFAULT_TIMEOUT_MS),
+    maxRetries: normalizeNumber(process.env.DROID_REFRESH_RETRIES, DEFAULT_MAX_RETRIES),
+    retryDelayMs: normalizeNumber(process.env.DROID_REFRESH_RETRY_BASE_MS, DEFAULT_RETRY_BASE_MS)
+  };
+}
+
+function sleep(ms) {
+  if (!ms || ms <= 0) return Promise.resolve();
+  return new Promise(resolve => setTimeout(resolve, ms));
+}
+
+function isRetryableError(error) {
+  if (!error) return false;
+  if (error.name === 'AbortError') return true;
+  const retryCodes = new Set(['ECONNRESET', 'ETIMEDOUT', 'ECONNREFUSED', 'EAI_AGAIN', 'ENOTFOUND']);
+  return retryCodes.has(error.code);
+}
+
+function shouldRetryStatus(status) {
+  return status === 429 || status >= 500;
+}
+
+function parseRetryAfterMs(response) {
+  if (!response?.headers?.get) return null;
+  const raw = response.headers.get('retry-after');
+  if (!raw) return null;
+
+  const seconds = parseInt(raw, 10);
+  if (Number.isFinite(seconds)) {
+    return Math.max(0, seconds * 1000);
+  }
+
+  const dateMs = Date.parse(raw);
+  if (!Number.isNaN(dateMs)) {
+    return Math.max(0, dateMs - Date.now());
+  }
+
+  return null;
+}
+
+async function fetchWithTimeout(url, options, timeoutMs, fetchImpl) {
+  const controller = new AbortController();
+  const timeoutId = setTimeout(() => controller.abort(), timeoutMs);
+  try {
+    return await fetchImpl(url, { ...options, signal: controller.signal });
+  } finally {
+    clearTimeout(timeoutId);
+  }
+}
+
+function buildError(status, errorText) {
+  const message = `Failed to refresh token: ${status} ${errorText || ''}`.trim();
+  const error = new Error(message);
+  error.status = status;
+  return error;
+}
+
+export async function requestRefreshToken(options) {
+  const {
+    refreshUrl,
+    refreshToken,
+    clientId,
+    proxyAgentInfo,
+    timeoutMs = DEFAULT_TIMEOUT_MS,
+    maxRetries = DEFAULT_MAX_RETRIES,
+    retryDelayMs = DEFAULT_RETRY_BASE_MS,
+    fetchImpl = fetch
+  } = options;
+
+  let attempt = 0;
+  while (true) {
+    try {
+      const formData = new URLSearchParams();
+      formData.append('grant_type', 'refresh_token');
+      formData.append('refresh_token', refreshToken);
+      formData.append('client_id', clientId);
+
+      const fetchOptions = {
+        method: 'POST',
+        headers: {
+          'Content-Type': 'application/x-www-form-urlencoded'
+        },
+        body: formData.toString()
+      };
+
+      if (proxyAgentInfo?.agent) {
+        fetchOptions.agent = proxyAgentInfo.agent;
+      }
+
+      const response = await fetchWithTimeout(refreshUrl, fetchOptions, timeoutMs, fetchImpl);
+
+      if (!response.ok) {
+        const errorText = await response.text().catch(() => '');
+        if (shouldRetryStatus(response.status) && attempt < maxRetries) {
+          const retryAfter = parseRetryAfterMs(response);
+          const delay = retryAfter ?? Math.min(retryDelayMs * (2 ** attempt), MAX_RETRY_DELAY_MS);
+          attempt += 1;
+          await sleep(delay);
+          continue;
+        }
+        throw buildError(response.status, errorText);
+      }
+
+      return await response.json();
+    } catch (error) {
+      if (isRetryableError(error) && attempt < maxRetries) {
+        const delay = Math.min(retryDelayMs * (2 ** attempt), MAX_RETRY_DELAY_MS);
+        attempt += 1;
+        await sleep(delay);
+        continue;
+      }
+      throw error;
+    }
+  }
+}
--- a/tests/refresh-client.test.js
+++ b/tests/refresh-client.test.js
@@ -0,0 +1,112 @@
+import test from 'node:test';
+import assert from 'node:assert/strict';
+import { requestRefreshToken } from '../refresh-client.js';
+
+function mockResponse({ ok, status, jsonData, textData, headers }) {
+  const normalizedHeaders = {};
+  if (headers) {
+    for (const [key, value] of Object.entries(headers)) {
+      normalizedHeaders[key.toLowerCase()] = value;
+    }
+  }
+
+  return {
+    ok,
+    status,
+    json: async () => jsonData ?? {},
+    text: async () => textData ?? '',
+    headers: {
+      get: (name) => normalizedHeaders[name.toLowerCase()] ?? null
+    }
+  };
+}
+
+test('refresh retries on 500 then succeeds', async () => {
+  let calls = 0;
+  const fetchImpl = async () => {
+    calls += 1;
+    if (calls === 1) {
+      return mockResponse({ ok: false, status: 500, textData: 'boom' });
+    }
+    return mockResponse({
+      ok: true,
+      status: 200,
+      jsonData: { access_token: 'access', refresh_token: 'refresh' }
+    });
+  };
+
+  const data = await requestRefreshToken({
+    refreshUrl: 'https://example.com',
+    refreshToken: 'refresh_token',
+    clientId: 'client',
+    timeoutMs: 20,
+    maxRetries: 1,
+    retryDelayMs: 1,
+    fetchImpl
+  });
+
+  assert.equal(data.access_token, 'access');
+  assert.equal(calls, 2);
+});
+
+test('refresh does not retry on 400', async () => {
+  let calls = 0;
+  const fetchImpl = async () => {
+    calls += 1;
+    return mockResponse({ ok: false, status: 400, textData: 'bad request' });
+  };
+
+  await assert.rejects(
+    () => requestRefreshToken({
+      refreshUrl: 'https://example.com',
+      refreshToken: 'refresh_token',
+      clientId: 'client',
+      timeoutMs: 20,
+      maxRetries: 2,
+      retryDelayMs: 1,
+      fetchImpl
+    }),
+    (err) => err?.status === 400
+  );
+
+  assert.equal(calls, 1);
+});
+
+test('refresh retries on timeout abort', async () => {
+  let calls = 0;
+  const fetchImpl = async (url, options) => {
+    calls += 1;
+    return new Promise((resolve, reject) => {
+      if (!options?.signal) {
+        resolve(mockResponse({ ok: true, status: 200, jsonData: {} }));
+        return;
+      }
+      if (options.signal.aborted) {
+        const error = new Error('Aborted');
+        error.name = 'AbortError';
+        reject(error);
+        return;
+      }
+      options.signal.addEventListener('abort', () => {
+        const error = new Error('Aborted');
+        error.name = 'AbortError';
+        reject(error);
+      });
+    });
+  };
+
+  await assert.rejects(
+    () => requestRefreshToken({
+      refreshUrl: 'https://example.com',
+      refreshToken: 'refresh_token',
+      clientId: 'client',
+      timeoutMs: 5,
+      maxRetries: 1,
+      retryDelayMs: 1,
+      fetchImpl
+    }),
+    (err) => err?.name === 'AbortError'
+  );
+
+  assert.equal(calls, 2);
+});