mirror of
https://github.com/Wei-Shaw/claude-relay-service.git
synced 2026-01-23 09:38:02 +00:00
feat: 添加多模型支持和OpenAI兼容接口
- 新增 Gemini 模型支持和账户管理功能 - 实现 OpenAI 格式到 Claude/Gemini 的请求转换 - 添加自动 token 刷新服务，支持提前刷新策略 - 增强 Web 管理界面，支持 Gemini 账户管理 - 优化 token 显示，添加掩码功能 - 完善日志记录和错误处理 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
shaw
@@ -17,6 +17,8 @@ class ApiKeyService {
|
||||
tokenLimit = config.limits.defaultTokenLimit,
|
||||
expiresAt = null,
|
||||
claudeAccountId = null,
|
||||
geminiAccountId = null,
|
||||
permissions = 'all', // 'claude', 'gemini', 'all'
|
||||
isActive = true,
|
||||
concurrencyLimit = 0,
|
||||
rateLimitWindow = null,
|
||||
@@ -41,6 +43,8 @@ class ApiKeyService {
|
||||
rateLimitRequests: String(rateLimitRequests ?? 0),
|
||||
isActive: String(isActive),
|
||||
claudeAccountId: claudeAccountId || '',
|
||||
geminiAccountId: geminiAccountId || '',
|
||||
permissions: permissions || 'all',
|
||||
enableModelRestriction: String(enableModelRestriction),
|
||||
restrictedModels: JSON.stringify(restrictedModels || []),
|
||||
createdAt: new Date().toISOString(),
|
||||
@@ -65,6 +69,8 @@ class ApiKeyService {
|
||||
rateLimitRequests: parseInt(keyData.rateLimitRequests || 0),
|
||||
isActive: keyData.isActive === 'true',
|
||||
claudeAccountId: keyData.claudeAccountId,
|
||||
geminiAccountId: keyData.geminiAccountId,
|
||||
permissions: keyData.permissions,
|
||||
enableModelRestriction: keyData.enableModelRestriction === 'true',
|
||||
restrictedModels: JSON.parse(keyData.restrictedModels),
|
||||
createdAt: keyData.createdAt,
|
||||
@@ -122,6 +128,8 @@ class ApiKeyService {
|
||||
id: keyData.id,
|
||||
name: keyData.name,
|
||||
claudeAccountId: keyData.claudeAccountId,
|
||||
geminiAccountId: keyData.geminiAccountId,
|
||||
permissions: keyData.permissions || 'all',
|
||||
tokenLimit: parseInt(keyData.tokenLimit),
|
||||
concurrencyLimit: parseInt(keyData.concurrencyLimit || 0),
|
||||
rateLimitWindow: parseInt(keyData.rateLimitWindow || 0),
|
||||
@@ -152,6 +160,7 @@ class ApiKeyService {
|
||||
key.currentConcurrency = await redis.getConcurrency(key.id);
|
||||
key.isActive = key.isActive === 'true';
|
||||
key.enableModelRestriction = key.enableModelRestriction === 'true';
|
||||
key.permissions = key.permissions || 'all'; // 兼容旧数据
|
||||
try {
|
||||
key.restrictedModels = key.restrictedModels ? JSON.parse(key.restrictedModels) : [];
|
||||
} catch (e) {
|
||||
@@ -176,7 +185,7 @@ class ApiKeyService {
|
||||
}
|
||||
|
||||
// 允许更新的字段
|
||||
const allowedUpdates = ['name', 'description', 'tokenLimit', 'concurrencyLimit', 'rateLimitWindow', 'rateLimitRequests', 'isActive', 'claudeAccountId', 'expiresAt', 'enableModelRestriction', 'restrictedModels'];
|
||||
const allowedUpdates = ['name', 'description', 'tokenLimit', 'concurrencyLimit', 'rateLimitWindow', 'rateLimitRequests', 'isActive', 'claudeAccountId', 'geminiAccountId', 'permissions', 'expiresAt', 'enableModelRestriction', 'restrictedModels'];
|
||||
const updatedData = { ...keyData };
|
||||
|
||||
for (const [field, value] of Object.entries(updates)) {
|
||||
@@ -292,4 +301,10 @@ class ApiKeyService {
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = new ApiKeyService();
|
||||
// 导出实例和单独的方法
|
||||
const apiKeyService = new ApiKeyService();
|
||||
|
||||
// 为了方便其他服务调用，导出 recordUsage 方法
|
||||
apiKeyService.recordUsageMetrics = apiKeyService.recordUsage.bind(apiKeyService);
|
||||
|
||||
module.exports = apiKeyService;
|
||||
@@ -6,6 +6,15 @@ const axios = require('axios');
|
||||
const redis = require('../models/redis');
|
||||
const logger = require('../utils/logger');
|
||||
const config = require('../../config/config');
|
||||
const { maskToken } = require('../utils/tokenMask');
|
||||
const {
|
||||
logRefreshStart,
|
||||
logRefreshSuccess,
|
||||
logRefreshError,
|
||||
logTokenUsage,
|
||||
logRefreshSkipped
|
||||
} = require('../utils/tokenRefreshLogger');
|
||||
const tokenRefreshService = require('./tokenRefreshService');
|
||||
|
||||
class ClaudeAccountService {
|
||||
constructor() {
|
||||
@@ -101,6 +110,8 @@ class ClaudeAccountService {
|
||||
|
||||
// 🔄 刷新Claude账户token
|
||||
async refreshAccountToken(accountId) {
|
||||
let lockAcquired = false;
|
||||
|
||||
try {
|
||||
const accountData = await redis.getClaudeAccount(accountId);
|
||||
|
||||
@@ -114,6 +125,35 @@ class ClaudeAccountService {
|
||||
throw new Error('No refresh token available - manual token update required');
|
||||
}
|
||||
|
||||
// 尝试获取分布式锁
|
||||
lockAcquired = await tokenRefreshService.acquireRefreshLock(accountId, 'claude');
|
||||
|
||||
if (!lockAcquired) {
|
||||
// 如果无法获取锁，说明另一个进程正在刷新
|
||||
logger.info(`🔒 Token refresh already in progress for account: ${accountData.name} (${accountId})`);
|
||||
logRefreshSkipped(accountId, accountData.name, 'claude', 'already_locked');
|
||||
|
||||
// 等待一段时间后返回，期望其他进程已完成刷新
|
||||
await new Promise(resolve => setTimeout(resolve, 2000));
|
||||
|
||||
// 重新获取账户数据（可能已被其他进程刷新）
|
||||
const updatedData = await redis.getClaudeAccount(accountId);
|
||||
if (updatedData && updatedData.accessToken) {
|
||||
const accessToken = this._decryptSensitiveData(updatedData.accessToken);
|
||||
return {
|
||||
success: true,
|
||||
accessToken: accessToken,
|
||||
expiresAt: updatedData.expiresAt
|
||||
};
|
||||
}
|
||||
|
||||
throw new Error('Token refresh in progress by another process');
|
||||
}
|
||||
|
||||
// 记录开始刷新
|
||||
logRefreshStart(accountId, accountData.name, 'claude', 'manual_refresh');
|
||||
logger.info(`🔄 Starting token refresh for account: ${accountData.name} (${accountId})`);
|
||||
|
||||
// 创建代理agent
|
||||
const agent = this._createProxyAgent(accountData.proxy);
|
||||
|
||||
@@ -125,7 +165,7 @@ class ClaudeAccountService {
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
'Accept': 'application/json, text/plain, */*',
|
||||
'User-Agent': 'claude-cli/1.0.53 (external, cli)',
|
||||
'User-Agent': 'claude-cli/1.0.56 (external, cli)',
|
||||
'Accept-Language': 'en-US,en;q=0.9',
|
||||
'Referer': 'https://claude.ai/',
|
||||
'Origin': 'https://claude.ai'
|
||||
@@ -147,7 +187,15 @@ class ClaudeAccountService {
|
||||
|
||||
await redis.setClaudeAccount(accountId, accountData);
|
||||
|
||||
logger.success(`🔄 Refreshed token for account: ${accountData.name} (${accountId})`);
|
||||
// 记录刷新成功
|
||||
logRefreshSuccess(accountId, accountData.name, 'claude', {
|
||||
accessToken: access_token,
|
||||
refreshToken: refresh_token,
|
||||
expiresAt: accountData.expiresAt,
|
||||
scopes: accountData.scopes
|
||||
});
|
||||
|
||||
logger.success(`🔄 Refreshed token for account: ${accountData.name} (${accountId}) - Access Token: ${maskToken(access_token)}`);
|
||||
|
||||
return {
|
||||
success: true,
|
||||
@@ -158,17 +206,23 @@ class ClaudeAccountService {
|
||||
throw new Error(`Token refresh failed with status: ${response.status}`);
|
||||
}
|
||||
} catch (error) {
|
||||
logger.error(`❌ Failed to refresh token for account ${accountId}:`, error);
|
||||
|
||||
// 更新错误状态
|
||||
// 记录刷新失败
|
||||
const accountData = await redis.getClaudeAccount(accountId);
|
||||
if (accountData) {
|
||||
logRefreshError(accountId, accountData.name, 'claude', error);
|
||||
accountData.status = 'error';
|
||||
accountData.errorMessage = error.message;
|
||||
await redis.setClaudeAccount(accountId, accountData);
|
||||
}
|
||||
|
||||
logger.error(`❌ Failed to refresh token for account ${accountId}:`, error);
|
||||
|
||||
throw error;
|
||||
} finally {
|
||||
// 释放锁
|
||||
if (lockAcquired) {
|
||||
await tokenRefreshService.releaseRefreshLock(accountId, 'claude');
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -188,8 +242,12 @@ class ClaudeAccountService {
|
||||
// 检查token是否过期
|
||||
const expiresAt = parseInt(accountData.expiresAt);
|
||||
const now = Date.now();
|
||||
const isExpired = !expiresAt || now >= (expiresAt - 60000); // 60秒提前刷新
|
||||
|
||||
if (!expiresAt || now >= (expiresAt - 60000)) { // 60秒提前刷新
|
||||
// 记录token使用情况
|
||||
logTokenUsage(accountId, accountData.name, 'claude', accountData.expiresAt, isExpired);
|
||||
|
||||
if (isExpired) {
|
||||
logger.info(`🔄 Token expired/expiring for account ${accountId}, attempting refresh...`);
|
||||
try {
|
||||
const refreshResult = await this.refreshAccountToken(accountId);
|
||||
@@ -275,6 +333,9 @@ class ClaudeAccountService {
|
||||
const allowedUpdates = ['name', 'description', 'email', 'password', 'refreshToken', 'proxy', 'isActive', 'claudeAiOauth', 'accountType'];
|
||||
const updatedData = { ...accountData };
|
||||
|
||||
// 检查是否新增了 refresh token
|
||||
const oldRefreshToken = this._decryptSensitiveData(accountData.refreshToken);
|
||||
|
||||
for (const [field, value] of Object.entries(updates)) {
|
||||
if (allowedUpdates.includes(field)) {
|
||||
if (['email', 'password', 'refreshToken'].includes(field)) {
|
||||
@@ -298,6 +359,27 @@ class ClaudeAccountService {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 如果新增了 refresh token（之前没有，现在有了），更新过期时间为10分钟
|
||||
if (updates.refreshToken && !oldRefreshToken && updates.refreshToken.trim()) {
|
||||
const newExpiresAt = Date.now() + (10 * 60 * 1000); // 10分钟
|
||||
updatedData.expiresAt = newExpiresAt.toString();
|
||||
logger.info(`🔄 New refresh token added for account ${accountId}, setting expiry to 10 minutes`);
|
||||
}
|
||||
|
||||
// 如果通过 claudeAiOauth 更新，也要检查是否新增了 refresh token
|
||||
if (updates.claudeAiOauth && updates.claudeAiOauth.refreshToken && !oldRefreshToken) {
|
||||
// 如果 expiresAt 设置的时间过长（超过1小时），调整为10分钟
|
||||
const providedExpiry = parseInt(updates.claudeAiOauth.expiresAt);
|
||||
const now = Date.now();
|
||||
const oneHour = 60 * 60 * 1000;
|
||||
|
||||
if (providedExpiry - now > oneHour) {
|
||||
const newExpiresAt = now + (10 * 60 * 1000); // 10分钟
|
||||
updatedData.expiresAt = newExpiresAt.toString();
|
||||
logger.info(`🔄 Adjusted expiry time to 10 minutes for account ${accountId} with refresh token`);
|
||||
}
|
||||
}
|
||||
|
||||
updatedData.updatedAt = new Date().toISOString();
|
||||
|
||||
|
||||
@@ -445,7 +445,7 @@ class ClaudeRelayService {
|
||||
}
|
||||
|
||||
// 🌊 处理流式响应（带usage数据捕获）
|
||||
async relayStreamRequestWithUsageCapture(requestBody, apiKeyData, responseStream, clientHeaders, usageCallback) {
|
||||
async relayStreamRequestWithUsageCapture(requestBody, apiKeyData, responseStream, clientHeaders, usageCallback, streamTransformer = null) {
|
||||
try {
|
||||
// 调试日志：查看API Key数据（流式请求）
|
||||
logger.info('🔍 [Stream] API Key data received:', {
|
||||
@@ -495,7 +495,7 @@ class ClaudeRelayService {
|
||||
const proxyAgent = await this._getProxyAgent(accountId);
|
||||
|
||||
// 发送流式请求并捕获usage数据
|
||||
return await this._makeClaudeStreamRequestWithUsageCapture(processedBody, accessToken, proxyAgent, clientHeaders, responseStream, usageCallback, accountId, sessionHash);
|
||||
return await this._makeClaudeStreamRequestWithUsageCapture(processedBody, accessToken, proxyAgent, clientHeaders, responseStream, usageCallback, accountId, sessionHash, streamTransformer);
|
||||
} catch (error) {
|
||||
logger.error('❌ Claude stream relay with usage capture failed:', error);
|
||||
throw error;
|
||||
@@ -503,7 +503,7 @@ class ClaudeRelayService {
|
||||
}
|
||||
|
||||
// 🌊 发送流式请求到Claude API（带usage数据捕获）
|
||||
async _makeClaudeStreamRequestWithUsageCapture(body, accessToken, proxyAgent, clientHeaders, responseStream, usageCallback, accountId, sessionHash) {
|
||||
async _makeClaudeStreamRequestWithUsageCapture(body, accessToken, proxyAgent, clientHeaders, responseStream, usageCallback, accountId, sessionHash, streamTransformer = null) {
|
||||
return new Promise((resolve, reject) => {
|
||||
const url = new URL(this.claudeApiUrl);
|
||||
|
||||
@@ -559,7 +559,15 @@ class ClaudeRelayService {
|
||||
// 转发已处理的完整行到客户端
|
||||
if (lines.length > 0) {
|
||||
const linesToForward = lines.join('\n') + (lines.length > 0 ? '\n' : '');
|
||||
responseStream.write(linesToForward);
|
||||
// 如果有流转换器，应用转换
|
||||
if (streamTransformer) {
|
||||
const transformed = streamTransformer(linesToForward);
|
||||
if (transformed) {
|
||||
responseStream.write(transformed);
|
||||
}
|
||||
} else {
|
||||
responseStream.write(linesToForward);
|
||||
}
|
||||
}
|
||||
|
||||
for (const line of lines) {
|
||||
@@ -612,7 +620,14 @@ class ClaudeRelayService {
|
||||
res.on('end', async () => {
|
||||
// 处理缓冲区中剩余的数据
|
||||
if (buffer.trim()) {
|
||||
responseStream.write(buffer);
|
||||
if (streamTransformer) {
|
||||
const transformed = streamTransformer(buffer);
|
||||
if (transformed) {
|
||||
responseStream.write(transformed);
|
||||
}
|
||||
} else {
|
||||
responseStream.write(buffer);
|
||||
}
|
||||
}
|
||||
responseStream.end();
|
||||
|
||||
|
||||
673
src/services/geminiAccountService.js
Normal file
673
src/services/geminiAccountService.js
Normal file
@@ -0,0 +1,673 @@
|
||||
const redisClient = require('../models/redis');
|
||||
const { v4: uuidv4 } = require('uuid');
|
||||
const crypto = require('crypto');
|
||||
const config = require('../../config/config');
|
||||
const logger = require('../utils/logger');
|
||||
const { OAuth2Client } = require('google-auth-library');
|
||||
const { maskToken } = require('../utils/tokenMask');
|
||||
const {
|
||||
logRefreshStart,
|
||||
logRefreshSuccess,
|
||||
logRefreshError,
|
||||
logTokenUsage,
|
||||
logRefreshSkipped
|
||||
} = require('../utils/tokenRefreshLogger');
|
||||
const tokenRefreshService = require('./tokenRefreshService');
|
||||
|
||||
// Gemini CLI OAuth 配置
|
||||
const OAUTH_CLIENT_ID = config.gemini.oauthClientId;
|
||||
const OAUTH_CLIENT_SECRET = config.gemini.oauthClientSecret;
|
||||
const OAUTH_SCOPES = ['https://www.googleapis.com/auth/cloud-platform'];
|
||||
|
||||
// 加密相关常量
|
||||
const ALGORITHM = 'aes-256-cbc';
|
||||
const ENCRYPTION_KEY = Buffer.from(config.security.encryptionKey, 'hex');
|
||||
const IV_LENGTH = 16;
|
||||
|
||||
// Gemini 账户键前缀
|
||||
const GEMINI_ACCOUNT_KEY_PREFIX = 'gemini_account:';
|
||||
const SHARED_GEMINI_ACCOUNTS_KEY = 'shared_gemini_accounts';
|
||||
const ACCOUNT_SESSION_MAPPING_PREFIX = 'gemini_session_account_mapping:';
|
||||
|
||||
// 加密函数
|
||||
function encrypt(text) {
|
||||
if (!text) return '';
|
||||
const iv = crypto.randomBytes(IV_LENGTH);
|
||||
const cipher = crypto.createCipheriv(ALGORITHM, ENCRYPTION_KEY, iv);
|
||||
let encrypted = cipher.update(text);
|
||||
encrypted = Buffer.concat([encrypted, cipher.final()]);
|
||||
return iv.toString('hex') + ':' + encrypted.toString('hex');
|
||||
}
|
||||
|
||||
// 解密函数
|
||||
function decrypt(text) {
|
||||
if (!text) return '';
|
||||
try {
|
||||
const textParts = text.split(':');
|
||||
const iv = Buffer.from(textParts.shift(), 'hex');
|
||||
const encryptedText = Buffer.from(textParts.join(':'), 'hex');
|
||||
const decipher = crypto.createDecipheriv(ALGORITHM, ENCRYPTION_KEY, iv);
|
||||
let decrypted = decipher.update(encryptedText);
|
||||
decrypted = Buffer.concat([decrypted, decipher.final()]);
|
||||
return decrypted.toString();
|
||||
} catch (error) {
|
||||
logger.error('Decryption error:', error);
|
||||
return '';
|
||||
}
|
||||
}
|
||||
|
||||
// 创建 OAuth2 客户端
|
||||
function createOAuth2Client(redirectUri = null) {
|
||||
// 如果没有提供 redirectUri，使用默认值
|
||||
const uri = redirectUri || 'http://localhost:8085';
|
||||
return new OAuth2Client(
|
||||
OAUTH_CLIENT_ID,
|
||||
OAUTH_CLIENT_SECRET,
|
||||
uri
|
||||
);
|
||||
}
|
||||
|
||||
// 生成授权 URL
|
||||
async function generateAuthUrl(state = null, redirectUri = null) {
|
||||
const oAuth2Client = createOAuth2Client(redirectUri);
|
||||
const authUrl = oAuth2Client.generateAuthUrl({
|
||||
access_type: 'offline',
|
||||
scope: OAUTH_SCOPES,
|
||||
prompt: 'select_account',
|
||||
state: state || uuidv4()
|
||||
});
|
||||
|
||||
return {
|
||||
authUrl,
|
||||
state: state || authUrl.split('state=')[1].split('&')[0]
|
||||
};
|
||||
}
|
||||
|
||||
// 轮询检查 OAuth 授权状态
|
||||
async function pollAuthorizationStatus(sessionId, maxAttempts = 60, interval = 2000) {
|
||||
let attempts = 0;
|
||||
const client = redisClient.getClientSafe();
|
||||
|
||||
while (attempts < maxAttempts) {
|
||||
try {
|
||||
const sessionData = await client.get(`oauth_session:${sessionId}`);
|
||||
if (!sessionData) {
|
||||
throw new Error('OAuth session not found');
|
||||
}
|
||||
|
||||
const session = JSON.parse(sessionData);
|
||||
if (session.code) {
|
||||
// 授权码已获取，交换 tokens
|
||||
const tokens = await exchangeCodeForTokens(session.code);
|
||||
|
||||
// 清理 session
|
||||
await client.del(`oauth_session:${sessionId}`);
|
||||
|
||||
return {
|
||||
success: true,
|
||||
tokens
|
||||
};
|
||||
}
|
||||
|
||||
if (session.error) {
|
||||
// 授权失败
|
||||
await client.del(`oauth_session:${sessionId}`);
|
||||
return {
|
||||
success: false,
|
||||
error: session.error
|
||||
};
|
||||
}
|
||||
|
||||
// 等待下一次轮询
|
||||
await new Promise(resolve => setTimeout(resolve, interval));
|
||||
attempts++;
|
||||
} catch (error) {
|
||||
logger.error('Error polling authorization status:', error);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
// 超时
|
||||
await client.del(`oauth_session:${sessionId}`);
|
||||
return {
|
||||
success: false,
|
||||
error: 'Authorization timeout'
|
||||
};
|
||||
}
|
||||
|
||||
// 交换授权码获取 tokens
|
||||
async function exchangeCodeForTokens(code, redirectUri = null) {
|
||||
const oAuth2Client = createOAuth2Client(redirectUri);
|
||||
|
||||
try {
|
||||
const { tokens } = await oAuth2Client.getToken(code);
|
||||
|
||||
// 转换为兼容格式
|
||||
return {
|
||||
access_token: tokens.access_token,
|
||||
refresh_token: tokens.refresh_token,
|
||||
scope: tokens.scope || OAUTH_SCOPES.join(' '),
|
||||
token_type: tokens.token_type || 'Bearer',
|
||||
expiry_date: tokens.expiry_date || Date.now() + (tokens.expires_in * 1000)
|
||||
};
|
||||
} catch (error) {
|
||||
logger.error('Error exchanging code for tokens:', error);
|
||||
throw new Error('Failed to exchange authorization code');
|
||||
}
|
||||
}
|
||||
|
||||
// 刷新访问令牌
|
||||
async function refreshAccessToken(refreshToken) {
|
||||
const oAuth2Client = createOAuth2Client();
|
||||
|
||||
try {
|
||||
oAuth2Client.setCredentials({
|
||||
refresh_token: refreshToken
|
||||
});
|
||||
|
||||
const { credentials } = await oAuth2Client.refreshAccessToken();
|
||||
|
||||
return {
|
||||
access_token: credentials.access_token,
|
||||
refresh_token: credentials.refresh_token || refreshToken,
|
||||
scope: credentials.scope || OAUTH_SCOPES.join(' '),
|
||||
token_type: credentials.token_type || 'Bearer',
|
||||
expiry_date: credentials.expiry_date
|
||||
};
|
||||
} catch (error) {
|
||||
logger.error('Error refreshing access token:', error);
|
||||
throw new Error('Failed to refresh access token');
|
||||
}
|
||||
}
|
||||
|
||||
// 创建 Gemini 账户
|
||||
async function createAccount(accountData) {
|
||||
const id = uuidv4();
|
||||
const now = new Date().toISOString();
|
||||
|
||||
// 处理凭证数据
|
||||
let geminiOauth = null;
|
||||
let accessToken = '';
|
||||
let refreshToken = '';
|
||||
let expiresAt = '';
|
||||
|
||||
if (accountData.geminiOauth || accountData.accessToken) {
|
||||
// 如果提供了完整的 OAuth 数据
|
||||
if (accountData.geminiOauth) {
|
||||
geminiOauth = typeof accountData.geminiOauth === 'string'
|
||||
? accountData.geminiOauth
|
||||
: JSON.stringify(accountData.geminiOauth);
|
||||
|
||||
const oauthData = typeof accountData.geminiOauth === 'string'
|
||||
? JSON.parse(accountData.geminiOauth)
|
||||
: accountData.geminiOauth;
|
||||
|
||||
accessToken = oauthData.access_token || '';
|
||||
refreshToken = oauthData.refresh_token || '';
|
||||
expiresAt = oauthData.expiry_date
|
||||
? new Date(oauthData.expiry_date).toISOString()
|
||||
: '';
|
||||
} else {
|
||||
// 如果只提供了 access token
|
||||
accessToken = accountData.accessToken;
|
||||
refreshToken = accountData.refreshToken || '';
|
||||
|
||||
// 构造完整的 OAuth 数据
|
||||
geminiOauth = JSON.stringify({
|
||||
access_token: accessToken,
|
||||
refresh_token: refreshToken,
|
||||
scope: accountData.scope || OAUTH_SCOPES.join(' '),
|
||||
token_type: accountData.tokenType || 'Bearer',
|
||||
expiry_date: accountData.expiryDate || Date.now() + 3600000 // 默认1小时
|
||||
});
|
||||
|
||||
expiresAt = new Date(accountData.expiryDate || Date.now() + 3600000).toISOString();
|
||||
}
|
||||
}
|
||||
|
||||
const account = {
|
||||
id,
|
||||
platform: 'gemini', // 标识为 Gemini 账户
|
||||
name: accountData.name || 'Gemini Account',
|
||||
description: accountData.description || '',
|
||||
accountType: accountData.accountType || 'shared',
|
||||
isActive: 'true',
|
||||
status: 'active',
|
||||
|
||||
// OAuth 相关字段（加密存储）
|
||||
geminiOauth: geminiOauth ? encrypt(geminiOauth) : '',
|
||||
accessToken: accessToken ? encrypt(accessToken) : '',
|
||||
refreshToken: refreshToken ? encrypt(refreshToken) : '',
|
||||
expiresAt,
|
||||
scopes: accountData.scopes || OAUTH_SCOPES.join(' '),
|
||||
|
||||
// 代理设置
|
||||
proxy: accountData.proxy ? JSON.stringify(accountData.proxy) : '',
|
||||
|
||||
// 项目编号（Google Cloud/Workspace 账号需要）
|
||||
projectId: accountData.projectId || '',
|
||||
|
||||
// 时间戳
|
||||
createdAt: now,
|
||||
updatedAt: now,
|
||||
lastUsedAt: '',
|
||||
lastRefreshAt: ''
|
||||
};
|
||||
|
||||
// 保存到 Redis
|
||||
const client = redisClient.getClientSafe();
|
||||
await client.hset(
|
||||
`${GEMINI_ACCOUNT_KEY_PREFIX}${id}`,
|
||||
account
|
||||
);
|
||||
|
||||
// 如果是共享账户，添加到共享账户集合
|
||||
if (account.accountType === 'shared') {
|
||||
await client.sadd(SHARED_GEMINI_ACCOUNTS_KEY, id);
|
||||
}
|
||||
|
||||
logger.info(`Created Gemini account: ${id}`);
|
||||
return account;
|
||||
}
|
||||
|
||||
// 获取账户
|
||||
async function getAccount(accountId) {
|
||||
const client = redisClient.getClientSafe();
|
||||
const accountData = await client.hgetall(`${GEMINI_ACCOUNT_KEY_PREFIX}${accountId}`);
|
||||
|
||||
if (!accountData || Object.keys(accountData).length === 0) {
|
||||
return null;
|
||||
}
|
||||
|
||||
// 解密敏感字段
|
||||
if (accountData.geminiOauth) {
|
||||
accountData.geminiOauth = decrypt(accountData.geminiOauth);
|
||||
}
|
||||
if (accountData.accessToken) {
|
||||
accountData.accessToken = decrypt(accountData.accessToken);
|
||||
}
|
||||
if (accountData.refreshToken) {
|
||||
accountData.refreshToken = decrypt(accountData.refreshToken);
|
||||
}
|
||||
|
||||
return accountData;
|
||||
}
|
||||
|
||||
// 更新账户
|
||||
async function updateAccount(accountId, updates) {
|
||||
const existingAccount = await getAccount(accountId);
|
||||
if (!existingAccount) {
|
||||
throw new Error('Account not found');
|
||||
}
|
||||
|
||||
const now = new Date().toISOString();
|
||||
updates.updatedAt = now;
|
||||
|
||||
// 检查是否新增了 refresh token
|
||||
const oldRefreshToken = existingAccount.refreshToken ? decrypt(existingAccount.refreshToken) : '';
|
||||
let needUpdateExpiry = false;
|
||||
|
||||
// 加密敏感字段
|
||||
if (updates.geminiOauth) {
|
||||
updates.geminiOauth = encrypt(
|
||||
typeof updates.geminiOauth === 'string'
|
||||
? updates.geminiOauth
|
||||
: JSON.stringify(updates.geminiOauth)
|
||||
);
|
||||
}
|
||||
if (updates.accessToken) {
|
||||
updates.accessToken = encrypt(updates.accessToken);
|
||||
}
|
||||
if (updates.refreshToken) {
|
||||
updates.refreshToken = encrypt(updates.refreshToken);
|
||||
// 如果之前没有 refresh token，现在有了，标记需要更新过期时间
|
||||
if (!oldRefreshToken && updates.refreshToken) {
|
||||
needUpdateExpiry = true;
|
||||
}
|
||||
}
|
||||
|
||||
// 更新账户类型时处理共享账户集合
|
||||
const client = redisClient.getClientSafe();
|
||||
if (updates.accountType && updates.accountType !== existingAccount.accountType) {
|
||||
if (updates.accountType === 'shared') {
|
||||
await client.sadd(SHARED_GEMINI_ACCOUNTS_KEY, accountId);
|
||||
} else {
|
||||
await client.srem(SHARED_GEMINI_ACCOUNTS_KEY, accountId);
|
||||
}
|
||||
}
|
||||
|
||||
// 如果新增了 refresh token，更新过期时间为10分钟
|
||||
if (needUpdateExpiry) {
|
||||
const newExpiry = new Date(Date.now() + (10 * 60 * 1000)).toISOString();
|
||||
updates.expiresAt = newExpiry;
|
||||
logger.info(`🔄 New refresh token added for Gemini account ${accountId}, setting expiry to 10 minutes`);
|
||||
}
|
||||
|
||||
// 如果通过 geminiOauth 更新，也要检查是否新增了 refresh token
|
||||
if (updates.geminiOauth && !oldRefreshToken) {
|
||||
const oauthData = typeof updates.geminiOauth === 'string'
|
||||
? JSON.parse(decrypt(updates.geminiOauth))
|
||||
: updates.geminiOauth;
|
||||
|
||||
if (oauthData.refresh_token) {
|
||||
// 如果 expiry_date 设置的时间过长（超过1小时），调整为10分钟
|
||||
const providedExpiry = oauthData.expiry_date || 0;
|
||||
const now = Date.now();
|
||||
const oneHour = 60 * 60 * 1000;
|
||||
|
||||
if (providedExpiry - now > oneHour) {
|
||||
const newExpiry = new Date(now + (10 * 60 * 1000)).toISOString();
|
||||
updates.expiresAt = newExpiry;
|
||||
logger.info(`🔄 Adjusted expiry time to 10 minutes for Gemini account ${accountId} with refresh token`);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
await client.hset(
|
||||
`${GEMINI_ACCOUNT_KEY_PREFIX}${accountId}`,
|
||||
updates
|
||||
);
|
||||
|
||||
logger.info(`Updated Gemini account: ${accountId}`);
|
||||
return { ...existingAccount, ...updates };
|
||||
}
|
||||
|
||||
// 删除账户
|
||||
async function deleteAccount(accountId) {
|
||||
const account = await getAccount(accountId);
|
||||
if (!account) {
|
||||
throw new Error('Account not found');
|
||||
}
|
||||
|
||||
// 从 Redis 删除
|
||||
const client = redisClient.getClientSafe();
|
||||
await client.del(`${GEMINI_ACCOUNT_KEY_PREFIX}${accountId}`);
|
||||
|
||||
// 从共享账户集合中移除
|
||||
if (account.accountType === 'shared') {
|
||||
await client.srem(SHARED_GEMINI_ACCOUNTS_KEY, accountId);
|
||||
}
|
||||
|
||||
// 清理会话映射
|
||||
const sessionMappings = await client.keys(`${ACCOUNT_SESSION_MAPPING_PREFIX}*`);
|
||||
for (const key of sessionMappings) {
|
||||
const mappedAccountId = await client.get(key);
|
||||
if (mappedAccountId === accountId) {
|
||||
await client.del(key);
|
||||
}
|
||||
}
|
||||
|
||||
logger.info(`Deleted Gemini account: ${accountId}`);
|
||||
return true;
|
||||
}
|
||||
|
||||
// 获取所有账户
|
||||
async function getAllAccounts() {
|
||||
const client = redisClient.getClientSafe();
|
||||
const keys = await client.keys(`${GEMINI_ACCOUNT_KEY_PREFIX}*`);
|
||||
const accounts = [];
|
||||
|
||||
for (const key of keys) {
|
||||
const accountData = await client.hgetall(key);
|
||||
if (accountData && Object.keys(accountData).length > 0) {
|
||||
// 不解密敏感字段，只返回基本信息
|
||||
accounts.push({
|
||||
...accountData,
|
||||
geminiOauth: accountData.geminiOauth ? '[ENCRYPTED]' : '',
|
||||
accessToken: accountData.accessToken ? '[ENCRYPTED]' : '',
|
||||
refreshToken: accountData.refreshToken ? '[ENCRYPTED]' : ''
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
return accounts;
|
||||
}
|
||||
|
||||
// 选择可用账户（支持专属和共享账户）
|
||||
async function selectAvailableAccount(apiKeyId, sessionHash = null) {
|
||||
// 首先检查是否有粘性会话
|
||||
const client = redisClient.getClientSafe();
|
||||
if (sessionHash) {
|
||||
const mappedAccountId = await client.get(
|
||||
`${ACCOUNT_SESSION_MAPPING_PREFIX}${sessionHash}`
|
||||
);
|
||||
|
||||
if (mappedAccountId) {
|
||||
const account = await getAccount(mappedAccountId);
|
||||
if (account && account.isActive === 'true' && !isTokenExpired(account)) {
|
||||
logger.debug(`Using sticky session account: ${mappedAccountId}`);
|
||||
return account;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 获取 API Key 信息
|
||||
const apiKeyData = await client.hgetall(`api_key:${apiKeyId}`);
|
||||
|
||||
// 检查是否绑定了 Gemini 账户
|
||||
if (apiKeyData.geminiAccountId) {
|
||||
const account = await getAccount(apiKeyData.geminiAccountId);
|
||||
if (account && account.isActive === 'true') {
|
||||
// 检查 token 是否过期
|
||||
const isExpired = isTokenExpired(account);
|
||||
|
||||
// 记录token使用情况
|
||||
logTokenUsage(account.id, account.name, 'gemini', account.expiresAt, isExpired);
|
||||
|
||||
if (isExpired) {
|
||||
await refreshAccountToken(account.id);
|
||||
return await getAccount(account.id);
|
||||
}
|
||||
|
||||
// 创建粘性会话映射
|
||||
if (sessionHash) {
|
||||
await client.setex(
|
||||
`${ACCOUNT_SESSION_MAPPING_PREFIX}${sessionHash}`,
|
||||
3600, // 1小时过期
|
||||
account.id
|
||||
);
|
||||
}
|
||||
|
||||
return account;
|
||||
}
|
||||
}
|
||||
|
||||
// 从共享账户池选择
|
||||
const sharedAccountIds = await client.smembers(SHARED_GEMINI_ACCOUNTS_KEY);
|
||||
const availableAccounts = [];
|
||||
|
||||
for (const accountId of sharedAccountIds) {
|
||||
const account = await getAccount(accountId);
|
||||
if (account && account.isActive === 'true' && !isRateLimited(account)) {
|
||||
availableAccounts.push(account);
|
||||
}
|
||||
}
|
||||
|
||||
if (availableAccounts.length === 0) {
|
||||
throw new Error('No available Gemini accounts');
|
||||
}
|
||||
|
||||
// 选择最少使用的账户
|
||||
availableAccounts.sort((a, b) => {
|
||||
const aLastUsed = a.lastUsedAt ? new Date(a.lastUsedAt).getTime() : 0;
|
||||
const bLastUsed = b.lastUsedAt ? new Date(b.lastUsedAt).getTime() : 0;
|
||||
return aLastUsed - bLastUsed;
|
||||
});
|
||||
|
||||
const selectedAccount = availableAccounts[0];
|
||||
|
||||
// 检查并刷新 token
|
||||
const isExpired = isTokenExpired(selectedAccount);
|
||||
|
||||
// 记录token使用情况
|
||||
logTokenUsage(selectedAccount.id, selectedAccount.name, 'gemini', selectedAccount.expiresAt, isExpired);
|
||||
|
||||
if (isExpired) {
|
||||
await refreshAccountToken(selectedAccount.id);
|
||||
return await getAccount(selectedAccount.id);
|
||||
}
|
||||
|
||||
// 创建粘性会话映射
|
||||
if (sessionHash) {
|
||||
await client.setex(
|
||||
`${ACCOUNT_SESSION_MAPPING_PREFIX}${sessionHash}`,
|
||||
3600,
|
||||
selectedAccount.id
|
||||
);
|
||||
}
|
||||
|
||||
return selectedAccount;
|
||||
}
|
||||
|
||||
// 检查 token 是否过期
|
||||
function isTokenExpired(account) {
|
||||
if (!account.expiresAt) return true;
|
||||
|
||||
const expiryTime = new Date(account.expiresAt).getTime();
|
||||
const now = Date.now();
|
||||
const buffer = 10 * 1000; // 10秒缓冲
|
||||
|
||||
return now >= (expiryTime - buffer);
|
||||
}
|
||||
|
||||
// 检查账户是否被限流
|
||||
function isRateLimited(account) {
|
||||
if (account.rateLimitStatus === 'limited' && account.rateLimitedAt) {
|
||||
const limitedAt = new Date(account.rateLimitedAt).getTime();
|
||||
const now = Date.now();
|
||||
const limitDuration = 60 * 60 * 1000; // 1小时
|
||||
|
||||
return now < (limitedAt + limitDuration);
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
// 刷新账户 token
|
||||
async function refreshAccountToken(accountId) {
|
||||
let lockAcquired = false;
|
||||
let account = null;
|
||||
|
||||
try {
|
||||
account = await getAccount(accountId);
|
||||
if (!account) {
|
||||
throw new Error('Account not found');
|
||||
}
|
||||
|
||||
if (!account.refreshToken) {
|
||||
throw new Error('No refresh token available');
|
||||
}
|
||||
|
||||
// 尝试获取分布式锁
|
||||
lockAcquired = await tokenRefreshService.acquireRefreshLock(accountId, 'gemini');
|
||||
|
||||
if (!lockAcquired) {
|
||||
// 如果无法获取锁，说明另一个进程正在刷新
|
||||
logger.info(`🔒 Token refresh already in progress for Gemini account: ${account.name} (${accountId})`);
|
||||
logRefreshSkipped(accountId, account.name, 'gemini', 'already_locked');
|
||||
|
||||
// 等待一段时间后返回，期望其他进程已完成刷新
|
||||
await new Promise(resolve => setTimeout(resolve, 2000));
|
||||
|
||||
// 重新获取账户数据（可能已被其他进程刷新）
|
||||
const updatedAccount = await getAccount(accountId);
|
||||
if (updatedAccount && updatedAccount.accessToken) {
|
||||
const accessToken = decrypt(updatedAccount.accessToken);
|
||||
return {
|
||||
access_token: accessToken,
|
||||
refresh_token: updatedAccount.refreshToken ? decrypt(updatedAccount.refreshToken) : '',
|
||||
expiry_date: updatedAccount.expiresAt ? new Date(updatedAccount.expiresAt).getTime() : 0,
|
||||
scope: updatedAccount.scope || OAUTH_SCOPES.join(' '),
|
||||
token_type: 'Bearer'
|
||||
};
|
||||
}
|
||||
|
||||
throw new Error('Token refresh in progress by another process');
|
||||
}
|
||||
|
||||
// 记录开始刷新
|
||||
logRefreshStart(accountId, account.name, 'gemini', 'manual_refresh');
|
||||
logger.info(`🔄 Starting token refresh for Gemini account: ${account.name} (${accountId})`);
|
||||
|
||||
const newTokens = await refreshAccessToken(decrypt(account.refreshToken));
|
||||
|
||||
// 更新账户信息
|
||||
const updates = {
|
||||
accessToken: newTokens.access_token,
|
||||
refreshToken: newTokens.refresh_token || account.refreshToken,
|
||||
expiresAt: new Date(newTokens.expiry_date).toISOString(),
|
||||
lastRefreshAt: new Date().toISOString(),
|
||||
geminiOauth: JSON.stringify(newTokens)
|
||||
};
|
||||
|
||||
await updateAccount(accountId, updates);
|
||||
|
||||
// 记录刷新成功
|
||||
logRefreshSuccess(accountId, account.name, 'gemini', {
|
||||
accessToken: newTokens.access_token,
|
||||
refreshToken: newTokens.refresh_token,
|
||||
expiresAt: newTokens.expiry_date,
|
||||
scopes: newTokens.scope
|
||||
});
|
||||
|
||||
logger.info(`Refreshed token for Gemini account: ${accountId} - Access Token: ${maskToken(newTokens.access_token)}`);
|
||||
|
||||
return newTokens;
|
||||
} catch (error) {
|
||||
// 记录刷新失败
|
||||
logRefreshError(accountId, account ? account.name : 'Unknown', 'gemini', error);
|
||||
|
||||
logger.error(`Failed to refresh token for account ${accountId}:`, error);
|
||||
|
||||
// 标记账户为错误状态
|
||||
await updateAccount(accountId, {
|
||||
status: 'error',
|
||||
errorMessage: error.message
|
||||
});
|
||||
|
||||
throw error;
|
||||
} finally {
|
||||
// 释放锁
|
||||
if (lockAcquired) {
|
||||
await tokenRefreshService.releaseRefreshLock(accountId, 'gemini');
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 标记账户被使用
|
||||
async function markAccountUsed(accountId) {
|
||||
await updateAccount(accountId, {
|
||||
lastUsedAt: new Date().toISOString()
|
||||
});
|
||||
}
|
||||
|
||||
// 设置账户限流状态
|
||||
async function setAccountRateLimited(accountId, isLimited = true) {
|
||||
const updates = isLimited ? {
|
||||
rateLimitStatus: 'limited',
|
||||
rateLimitedAt: new Date().toISOString()
|
||||
} : {
|
||||
rateLimitStatus: '',
|
||||
rateLimitedAt: ''
|
||||
};
|
||||
|
||||
await updateAccount(accountId, updates);
|
||||
}
|
||||
|
||||
module.exports = {
|
||||
generateAuthUrl,
|
||||
pollAuthorizationStatus,
|
||||
exchangeCodeForTokens,
|
||||
refreshAccessToken,
|
||||
createAccount,
|
||||
getAccount,
|
||||
updateAccount,
|
||||
deleteAccount,
|
||||
getAllAccounts,
|
||||
selectAvailableAccount,
|
||||
refreshAccountToken,
|
||||
markAccountUsed,
|
||||
setAccountRateLimited,
|
||||
isTokenExpired,
|
||||
OAUTH_CLIENT_ID,
|
||||
OAUTH_SCOPES
|
||||
};
|
||||
379
src/services/geminiRelayService.js
Normal file
379
src/services/geminiRelayService.js
Normal file
@@ -0,0 +1,379 @@
|
||||
const axios = require('axios');
|
||||
const { HttpsProxyAgent } = require('https-proxy-agent');
|
||||
const { SocksProxyAgent } = require('socks-proxy-agent');
|
||||
const logger = require('../utils/logger');
|
||||
const config = require('../../config/config');
|
||||
const { recordUsageMetrics } = require('./apiKeyService');
|
||||
|
||||
// Gemini API 配置
|
||||
const GEMINI_API_BASE = 'https://cloudcode.googleapis.com/v1';
|
||||
const DEFAULT_MODEL = 'models/gemini-2.0-flash-exp';
|
||||
|
||||
// 创建代理 agent
|
||||
function createProxyAgent(proxyConfig) {
|
||||
if (!proxyConfig) return null;
|
||||
|
||||
try {
|
||||
const proxy = typeof proxyConfig === 'string' ? JSON.parse(proxyConfig) : proxyConfig;
|
||||
|
||||
if (!proxy.type || !proxy.host || !proxy.port) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const proxyUrl = proxy.username && proxy.password
|
||||
? `${proxy.type}://${proxy.username}:${proxy.password}@${proxy.host}:${proxy.port}`
|
||||
: `${proxy.type}://${proxy.host}:${proxy.port}`;
|
||||
|
||||
if (proxy.type === 'socks5') {
|
||||
return new SocksProxyAgent(proxyUrl);
|
||||
} else if (proxy.type === 'http' || proxy.type === 'https') {
|
||||
return new HttpsProxyAgent(proxyUrl);
|
||||
}
|
||||
} catch (error) {
|
||||
logger.error('Error creating proxy agent:', error);
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
// 转换 OpenAI 消息格式到 Gemini 格式
|
||||
function convertMessagesToGemini(messages) {
|
||||
const contents = [];
|
||||
let systemInstruction = '';
|
||||
|
||||
for (const message of messages) {
|
||||
if (message.role === 'system') {
|
||||
systemInstruction += (systemInstruction ? '\n\n' : '') + message.content;
|
||||
} else if (message.role === 'user') {
|
||||
contents.push({
|
||||
role: 'user',
|
||||
parts: [{ text: message.content }]
|
||||
});
|
||||
} else if (message.role === 'assistant') {
|
||||
contents.push({
|
||||
role: 'model',
|
||||
parts: [{ text: message.content }]
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
return { contents, systemInstruction };
|
||||
}
|
||||
|
||||
// 转换 Gemini 响应到 OpenAI 格式
|
||||
function convertGeminiResponse(geminiResponse, model, stream = false) {
|
||||
if (stream) {
|
||||
// 流式响应
|
||||
const candidate = geminiResponse.candidates?.[0];
|
||||
if (!candidate) return null;
|
||||
|
||||
const content = candidate.content?.parts?.[0]?.text || '';
|
||||
const finishReason = candidate.finishReason?.toLowerCase();
|
||||
|
||||
return {
|
||||
id: `chatcmpl-${Date.now()}`,
|
||||
object: 'chat.completion.chunk',
|
||||
created: Math.floor(Date.now() / 1000),
|
||||
model: model,
|
||||
choices: [{
|
||||
index: 0,
|
||||
delta: {
|
||||
content: content
|
||||
},
|
||||
finish_reason: finishReason === 'stop' ? 'stop' : null
|
||||
}]
|
||||
};
|
||||
} else {
|
||||
// 非流式响应
|
||||
const candidate = geminiResponse.candidates?.[0];
|
||||
if (!candidate) {
|
||||
throw new Error('No response from Gemini');
|
||||
}
|
||||
|
||||
const content = candidate.content?.parts?.[0]?.text || '';
|
||||
const finishReason = candidate.finishReason?.toLowerCase() || 'stop';
|
||||
|
||||
// 计算 token 使用量
|
||||
const usage = geminiResponse.usageMetadata || {
|
||||
promptTokenCount: 0,
|
||||
candidatesTokenCount: 0,
|
||||
totalTokenCount: 0
|
||||
};
|
||||
|
||||
return {
|
||||
id: `chatcmpl-${Date.now()}`,
|
||||
object: 'chat.completion',
|
||||
created: Math.floor(Date.now() / 1000),
|
||||
model: model,
|
||||
choices: [{
|
||||
index: 0,
|
||||
message: {
|
||||
role: 'assistant',
|
||||
content: content
|
||||
},
|
||||
finish_reason: finishReason
|
||||
}],
|
||||
usage: {
|
||||
prompt_tokens: usage.promptTokenCount,
|
||||
completion_tokens: usage.candidatesTokenCount,
|
||||
total_tokens: usage.totalTokenCount
|
||||
}
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// 处理流式响应
|
||||
async function* handleStreamResponse(response, model, apiKeyId) {
|
||||
let buffer = '';
|
||||
let totalUsage = {
|
||||
promptTokenCount: 0,
|
||||
candidatesTokenCount: 0,
|
||||
totalTokenCount: 0
|
||||
};
|
||||
|
||||
try {
|
||||
for await (const chunk of response.data) {
|
||||
buffer += chunk.toString();
|
||||
|
||||
// 处理可能的多个 JSON 对象
|
||||
const lines = buffer.split('\n');
|
||||
buffer = lines.pop() || ''; // 保留最后一个不完整的行
|
||||
|
||||
for (const line of lines) {
|
||||
if (!line.trim()) continue;
|
||||
|
||||
try {
|
||||
const data = JSON.parse(line);
|
||||
|
||||
// 更新使用量统计
|
||||
if (data.usageMetadata) {
|
||||
totalUsage = data.usageMetadata;
|
||||
}
|
||||
|
||||
// 转换并发送响应
|
||||
const openaiResponse = convertGeminiResponse(data, model, true);
|
||||
if (openaiResponse) {
|
||||
yield `data: ${JSON.stringify(openaiResponse)}\n\n`;
|
||||
}
|
||||
|
||||
// 检查是否结束
|
||||
if (data.candidates?.[0]?.finishReason === 'STOP') {
|
||||
// 记录使用量
|
||||
if (apiKeyId && totalUsage.totalTokenCount > 0) {
|
||||
await recordUsageMetrics(apiKeyId, {
|
||||
inputTokens: totalUsage.promptTokenCount,
|
||||
outputTokens: totalUsage.candidatesTokenCount,
|
||||
model: model
|
||||
});
|
||||
}
|
||||
|
||||
yield 'data: [DONE]\n\n';
|
||||
return;
|
||||
}
|
||||
} catch (e) {
|
||||
logger.debug('Error parsing JSON line:', e.message);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 处理剩余的 buffer
|
||||
if (buffer.trim()) {
|
||||
try {
|
||||
const data = JSON.parse(buffer);
|
||||
const openaiResponse = convertGeminiResponse(data, model, true);
|
||||
if (openaiResponse) {
|
||||
yield `data: ${JSON.stringify(openaiResponse)}\n\n`;
|
||||
}
|
||||
} catch (e) {
|
||||
logger.debug('Error parsing final buffer:', e.message);
|
||||
}
|
||||
}
|
||||
|
||||
yield 'data: [DONE]\n\n';
|
||||
} catch (error) {
|
||||
logger.error('Stream processing error:', error);
|
||||
yield `data: ${JSON.stringify({
|
||||
error: {
|
||||
message: error.message,
|
||||
type: 'stream_error'
|
||||
}
|
||||
})}\n\n`;
|
||||
}
|
||||
}
|
||||
|
||||
// 发送请求到 Gemini
|
||||
async function sendGeminiRequest({
|
||||
messages,
|
||||
model = DEFAULT_MODEL,
|
||||
temperature = 0.7,
|
||||
maxTokens = 4096,
|
||||
stream = false,
|
||||
accessToken,
|
||||
proxy,
|
||||
apiKeyId,
|
||||
projectId,
|
||||
location = 'us-central1'
|
||||
}) {
|
||||
// 确保模型名称格式正确
|
||||
if (!model.startsWith('models/')) {
|
||||
model = `models/${model}`;
|
||||
}
|
||||
|
||||
// 转换消息格式
|
||||
const { contents, systemInstruction } = convertMessagesToGemini(messages);
|
||||
|
||||
// 构建请求体
|
||||
const requestBody = {
|
||||
contents,
|
||||
generationConfig: {
|
||||
temperature,
|
||||
maxOutputTokens: maxTokens,
|
||||
candidateCount: 1
|
||||
}
|
||||
};
|
||||
|
||||
if (systemInstruction) {
|
||||
requestBody.systemInstruction = { parts: [{ text: systemInstruction }] };
|
||||
}
|
||||
|
||||
// 配置请求选项
|
||||
let apiUrl;
|
||||
if (projectId) {
|
||||
// 使用项目特定的 URL 格式（Google Cloud/Workspace 账号）
|
||||
apiUrl = `${GEMINI_API_BASE}/projects/${projectId}/locations/${location}/${model}:${stream ? 'streamGenerateContent' : 'generateContent'}?alt=sse`;
|
||||
logger.debug(`Using project-specific URL with projectId: ${projectId}, location: ${location}`);
|
||||
} else {
|
||||
// 使用标准 URL 格式（个人 Google 账号）
|
||||
apiUrl = `${GEMINI_API_BASE}/${model}:${stream ? 'streamGenerateContent' : 'generateContent'}?alt=sse`;
|
||||
logger.debug('Using standard URL without projectId');
|
||||
}
|
||||
|
||||
const axiosConfig = {
|
||||
method: 'POST',
|
||||
url: apiUrl,
|
||||
headers: {
|
||||
'Authorization': `Bearer ${accessToken}`,
|
||||
'Content-Type': 'application/json'
|
||||
},
|
||||
data: requestBody,
|
||||
timeout: config.requestTimeout || 120000
|
||||
};
|
||||
|
||||
// 添加代理配置
|
||||
const proxyAgent = createProxyAgent(proxy);
|
||||
if (proxyAgent) {
|
||||
axiosConfig.httpsAgent = proxyAgent;
|
||||
logger.debug('Using proxy for Gemini request');
|
||||
}
|
||||
|
||||
if (stream) {
|
||||
axiosConfig.responseType = 'stream';
|
||||
}
|
||||
|
||||
try {
|
||||
logger.debug('Sending request to Gemini API');
|
||||
const response = await axios(axiosConfig);
|
||||
|
||||
if (stream) {
|
||||
return handleStreamResponse(response, model, apiKeyId);
|
||||
} else {
|
||||
// 非流式响应
|
||||
const openaiResponse = convertGeminiResponse(response.data, model, false);
|
||||
|
||||
// 记录使用量
|
||||
if (apiKeyId && openaiResponse.usage) {
|
||||
await recordUsageMetrics(apiKeyId, {
|
||||
inputTokens: openaiResponse.usage.prompt_tokens,
|
||||
outputTokens: openaiResponse.usage.completion_tokens,
|
||||
model: model
|
||||
});
|
||||
}
|
||||
|
||||
return openaiResponse;
|
||||
}
|
||||
} catch (error) {
|
||||
logger.error('Gemini API request failed:', error.response?.data || error.message);
|
||||
|
||||
// 转换错误格式
|
||||
if (error.response) {
|
||||
const geminiError = error.response.data?.error;
|
||||
throw {
|
||||
status: error.response.status,
|
||||
error: {
|
||||
message: geminiError?.message || 'Gemini API request failed',
|
||||
type: geminiError?.code || 'api_error',
|
||||
code: geminiError?.code
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
throw {
|
||||
status: 500,
|
||||
error: {
|
||||
message: error.message,
|
||||
type: 'network_error'
|
||||
}
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// 获取可用模型列表
|
||||
async function getAvailableModels(accessToken, proxy, projectId, location = 'us-central1') {
|
||||
let apiUrl;
|
||||
if (projectId) {
|
||||
// 使用项目特定的 URL 格式
|
||||
apiUrl = `${GEMINI_API_BASE}/projects/${projectId}/locations/${location}/models`;
|
||||
logger.debug(`Fetching models with projectId: ${projectId}, location: ${location}`);
|
||||
} else {
|
||||
// 使用标准 URL 格式
|
||||
apiUrl = `${GEMINI_API_BASE}/models`;
|
||||
logger.debug('Fetching models without projectId');
|
||||
}
|
||||
|
||||
const axiosConfig = {
|
||||
method: 'GET',
|
||||
url: apiUrl,
|
||||
headers: {
|
||||
'Authorization': `Bearer ${accessToken}`
|
||||
},
|
||||
timeout: 30000
|
||||
};
|
||||
|
||||
const proxyAgent = createProxyAgent(proxy);
|
||||
if (proxyAgent) {
|
||||
axiosConfig.httpsAgent = proxyAgent;
|
||||
}
|
||||
|
||||
try {
|
||||
const response = await axios(axiosConfig);
|
||||
const models = response.data.models || [];
|
||||
|
||||
// 转换为 OpenAI 格式
|
||||
return models
|
||||
.filter(model => model.supportedGenerationMethods?.includes('generateContent'))
|
||||
.map(model => ({
|
||||
id: model.name.replace('models/', ''),
|
||||
object: 'model',
|
||||
created: Date.now() / 1000,
|
||||
owned_by: 'google'
|
||||
}));
|
||||
} catch (error) {
|
||||
logger.error('Failed to get Gemini models:', error);
|
||||
// 返回默认模型列表
|
||||
return [
|
||||
{
|
||||
id: 'gemini-2.0-flash-exp',
|
||||
object: 'model',
|
||||
created: Date.now() / 1000,
|
||||
owned_by: 'google'
|
||||
}
|
||||
];
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = {
|
||||
sendGeminiRequest,
|
||||
getAvailableModels,
|
||||
convertMessagesToGemini,
|
||||
convertGeminiResponse
|
||||
};
|
||||
381
src/services/openaiToClaude.js
Normal file
381
src/services/openaiToClaude.js
Normal file
@@ -0,0 +1,381 @@
|
||||
/**
|
||||
* OpenAI 到 Claude 格式转换服务
|
||||
* 处理 OpenAI API 格式与 Claude API 格式之间的转换
|
||||
*/
|
||||
|
||||
const logger = require('../utils/logger');
|
||||
|
||||
class OpenAIToClaudeConverter {
|
||||
constructor() {
|
||||
// 停止原因映射
|
||||
this.stopReasonMapping = {
|
||||
'end_turn': 'stop',
|
||||
'max_tokens': 'length',
|
||||
'stop_sequence': 'stop',
|
||||
'tool_use': 'tool_calls'
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* 将 OpenAI 请求格式转换为 Claude 格式
|
||||
* @param {Object} openaiRequest - OpenAI 格式的请求
|
||||
* @returns {Object} Claude 格式的请求
|
||||
*/
|
||||
convertRequest(openaiRequest) {
|
||||
const claudeRequest = {
|
||||
model: openaiRequest.model, // 直接使用提供的模型名，不进行映射
|
||||
messages: this._convertMessages(openaiRequest.messages),
|
||||
max_tokens: openaiRequest.max_tokens || 4096,
|
||||
temperature: openaiRequest.temperature,
|
||||
top_p: openaiRequest.top_p,
|
||||
stream: openaiRequest.stream || false
|
||||
};
|
||||
|
||||
// 处理系统消息
|
||||
const systemMessage = this._extractSystemMessage(openaiRequest.messages);
|
||||
if (systemMessage) {
|
||||
claudeRequest.system = systemMessage;
|
||||
}
|
||||
|
||||
// 处理停止序列
|
||||
if (openaiRequest.stop) {
|
||||
claudeRequest.stop_sequences = Array.isArray(openaiRequest.stop)
|
||||
? openaiRequest.stop
|
||||
: [openaiRequest.stop];
|
||||
}
|
||||
|
||||
// 处理工具调用
|
||||
if (openaiRequest.tools) {
|
||||
claudeRequest.tools = this._convertTools(openaiRequest.tools);
|
||||
if (openaiRequest.tool_choice) {
|
||||
claudeRequest.tool_choice = this._convertToolChoice(openaiRequest.tool_choice);
|
||||
}
|
||||
}
|
||||
|
||||
// OpenAI 特有的参数已在转换过程中被忽略
|
||||
// 包括: n, presence_penalty, frequency_penalty, logit_bias, user
|
||||
|
||||
logger.debug('📝 Converted OpenAI request to Claude format:', {
|
||||
model: claudeRequest.model,
|
||||
messageCount: claudeRequest.messages.length,
|
||||
hasSystem: !!claudeRequest.system,
|
||||
stream: claudeRequest.stream
|
||||
});
|
||||
|
||||
return claudeRequest;
|
||||
}
|
||||
|
||||
/**
|
||||
* 将 Claude 响应格式转换为 OpenAI 格式
|
||||
* @param {Object} claudeResponse - Claude 格式的响应
|
||||
* @param {String} requestModel - 原始请求的模型名
|
||||
* @returns {Object} OpenAI 格式的响应
|
||||
*/
|
||||
convertResponse(claudeResponse, requestModel) {
|
||||
const timestamp = Math.floor(Date.now() / 1000);
|
||||
|
||||
const openaiResponse = {
|
||||
id: `chatcmpl-${this._generateId()}`,
|
||||
object: 'chat.completion',
|
||||
created: timestamp,
|
||||
model: requestModel || 'gpt-4',
|
||||
choices: [{
|
||||
index: 0,
|
||||
message: this._convertClaudeMessage(claudeResponse),
|
||||
finish_reason: this._mapStopReason(claudeResponse.stop_reason)
|
||||
}],
|
||||
usage: this._convertUsage(claudeResponse.usage)
|
||||
};
|
||||
|
||||
logger.debug('📝 Converted Claude response to OpenAI format:', {
|
||||
responseId: openaiResponse.id,
|
||||
finishReason: openaiResponse.choices[0].finish_reason,
|
||||
usage: openaiResponse.usage
|
||||
});
|
||||
|
||||
return openaiResponse;
|
||||
}
|
||||
|
||||
/**
|
||||
* 转换流式响应的单个数据块
|
||||
* @param {String} chunk - Claude SSE 数据块
|
||||
* @param {String} requestModel - 原始请求的模型名
|
||||
* @returns {String} OpenAI 格式的 SSE 数据块
|
||||
*/
|
||||
convertStreamChunk(chunk, requestModel) {
|
||||
if (!chunk || chunk.trim() === '') return '';
|
||||
|
||||
// 解析 SSE 数据
|
||||
const lines = chunk.split('\n');
|
||||
let convertedChunks = [];
|
||||
|
||||
for (const line of lines) {
|
||||
if (line.startsWith('data: ')) {
|
||||
const data = line.substring(6);
|
||||
if (data === '[DONE]') {
|
||||
convertedChunks.push('data: [DONE]\n\n');
|
||||
continue;
|
||||
}
|
||||
|
||||
try {
|
||||
const claudeEvent = JSON.parse(data);
|
||||
const openaiChunk = this._convertStreamEvent(claudeEvent, requestModel);
|
||||
if (openaiChunk) {
|
||||
convertedChunks.push(`data: ${JSON.stringify(openaiChunk)}\n\n`);
|
||||
}
|
||||
} catch (e) {
|
||||
// 如果不是 JSON，原样传递
|
||||
convertedChunks.push(line + '\n');
|
||||
}
|
||||
} else if (line.startsWith('event:') || line === '') {
|
||||
// 保留事件类型行和空行
|
||||
convertedChunks.push(line + '\n');
|
||||
}
|
||||
}
|
||||
|
||||
return convertedChunks.join('');
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 提取系统消息
|
||||
*/
|
||||
_extractSystemMessage(messages) {
|
||||
const systemMessages = messages.filter(msg => msg.role === 'system');
|
||||
if (systemMessages.length === 0) return null;
|
||||
|
||||
// 合并所有系统消息
|
||||
return systemMessages.map(msg => msg.content).join('\n\n');
|
||||
}
|
||||
|
||||
/**
|
||||
* 转换消息格式
|
||||
*/
|
||||
_convertMessages(messages) {
|
||||
const claudeMessages = [];
|
||||
|
||||
for (const msg of messages) {
|
||||
// 跳过系统消息（已经在 system 字段处理）
|
||||
if (msg.role === 'system') continue;
|
||||
|
||||
// 转换角色名称
|
||||
const role = msg.role === 'user' ? 'user' : 'assistant';
|
||||
|
||||
// 转换消息内容
|
||||
let content;
|
||||
if (typeof msg.content === 'string') {
|
||||
content = msg.content;
|
||||
} else if (Array.isArray(msg.content)) {
|
||||
// 处理多模态内容
|
||||
content = this._convertMultimodalContent(msg.content);
|
||||
} else {
|
||||
content = JSON.stringify(msg.content);
|
||||
}
|
||||
|
||||
const claudeMsg = {
|
||||
role: role,
|
||||
content: content
|
||||
};
|
||||
|
||||
// 处理工具调用
|
||||
if (msg.tool_calls) {
|
||||
claudeMsg.content = this._convertToolCalls(msg.tool_calls);
|
||||
}
|
||||
|
||||
// 处理工具响应
|
||||
if (msg.role === 'tool') {
|
||||
claudeMsg.role = 'user';
|
||||
claudeMsg.content = [{
|
||||
type: 'tool_result',
|
||||
tool_use_id: msg.tool_call_id,
|
||||
content: msg.content
|
||||
}];
|
||||
}
|
||||
|
||||
claudeMessages.push(claudeMsg);
|
||||
}
|
||||
|
||||
return claudeMessages;
|
||||
}
|
||||
|
||||
/**
|
||||
* 转换多模态内容
|
||||
*/
|
||||
_convertMultimodalContent(content) {
|
||||
return content.map(item => {
|
||||
if (item.type === 'text') {
|
||||
return {
|
||||
type: 'text',
|
||||
text: item.text
|
||||
};
|
||||
} else if (item.type === 'image_url') {
|
||||
return {
|
||||
type: 'image',
|
||||
source: {
|
||||
type: 'base64',
|
||||
media_type: 'image/jpeg', // 默认类型
|
||||
data: item.image_url.url.split(',')[1] // 假设是 base64
|
||||
}
|
||||
};
|
||||
}
|
||||
return item;
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* 转换工具定义
|
||||
*/
|
||||
_convertTools(tools) {
|
||||
return tools.map(tool => {
|
||||
if (tool.type === 'function') {
|
||||
return {
|
||||
name: tool.function.name,
|
||||
description: tool.function.description,
|
||||
input_schema: tool.function.parameters
|
||||
};
|
||||
}
|
||||
return tool;
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* 转换工具选择
|
||||
*/
|
||||
_convertToolChoice(toolChoice) {
|
||||
if (toolChoice === 'none') return { type: 'none' };
|
||||
if (toolChoice === 'auto') return { type: 'auto' };
|
||||
if (toolChoice === 'required') return { type: 'any' };
|
||||
if (toolChoice.type === 'function') {
|
||||
return {
|
||||
type: 'tool',
|
||||
name: toolChoice.function.name
|
||||
};
|
||||
}
|
||||
return { type: 'auto' };
|
||||
}
|
||||
|
||||
/**
|
||||
* 转换工具调用
|
||||
*/
|
||||
_convertToolCalls(toolCalls) {
|
||||
return toolCalls.map(tc => ({
|
||||
type: 'tool_use',
|
||||
id: tc.id,
|
||||
name: tc.function.name,
|
||||
input: JSON.parse(tc.function.arguments)
|
||||
}));
|
||||
}
|
||||
|
||||
/**
|
||||
* 转换 Claude 消息为 OpenAI 格式
|
||||
*/
|
||||
_convertClaudeMessage(claudeResponse) {
|
||||
const message = {
|
||||
role: 'assistant',
|
||||
content: null
|
||||
};
|
||||
|
||||
// 处理内容
|
||||
if (claudeResponse.content) {
|
||||
if (typeof claudeResponse.content === 'string') {
|
||||
message.content = claudeResponse.content;
|
||||
} else if (Array.isArray(claudeResponse.content)) {
|
||||
// 提取文本内容和工具调用
|
||||
const textParts = [];
|
||||
const toolCalls = [];
|
||||
|
||||
for (const item of claudeResponse.content) {
|
||||
if (item.type === 'text') {
|
||||
textParts.push(item.text);
|
||||
} else if (item.type === 'tool_use') {
|
||||
toolCalls.push({
|
||||
id: item.id,
|
||||
type: 'function',
|
||||
function: {
|
||||
name: item.name,
|
||||
arguments: JSON.stringify(item.input)
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
message.content = textParts.join('') || null;
|
||||
if (toolCalls.length > 0) {
|
||||
message.tool_calls = toolCalls;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return message;
|
||||
}
|
||||
|
||||
/**
|
||||
* 转换停止原因
|
||||
*/
|
||||
_mapStopReason(claudeReason) {
|
||||
return this.stopReasonMapping[claudeReason] || 'stop';
|
||||
}
|
||||
|
||||
/**
|
||||
* 转换使用统计
|
||||
*/
|
||||
_convertUsage(claudeUsage) {
|
||||
if (!claudeUsage) return undefined;
|
||||
|
||||
return {
|
||||
prompt_tokens: claudeUsage.input_tokens || 0,
|
||||
completion_tokens: claudeUsage.output_tokens || 0,
|
||||
total_tokens: (claudeUsage.input_tokens || 0) + (claudeUsage.output_tokens || 0)
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* 转换流式事件
|
||||
*/
|
||||
_convertStreamEvent(event, requestModel) {
|
||||
const timestamp = Math.floor(Date.now() / 1000);
|
||||
const baseChunk = {
|
||||
id: `chatcmpl-${this._generateId()}`,
|
||||
object: 'chat.completion.chunk',
|
||||
created: timestamp,
|
||||
model: requestModel || 'gpt-4',
|
||||
choices: [{
|
||||
index: 0,
|
||||
delta: {},
|
||||
finish_reason: null
|
||||
}]
|
||||
};
|
||||
|
||||
// 根据事件类型处理
|
||||
if (event.type === 'content_block_start' && event.content_block) {
|
||||
if (event.content_block.type === 'text') {
|
||||
baseChunk.choices[0].delta.content = event.content_block.text || '';
|
||||
}
|
||||
} else if (event.type === 'content_block_delta' && event.delta) {
|
||||
if (event.delta.type === 'text_delta') {
|
||||
baseChunk.choices[0].delta.content = event.delta.text || '';
|
||||
}
|
||||
} else if (event.type === 'message_delta' && event.delta) {
|
||||
if (event.delta.stop_reason) {
|
||||
baseChunk.choices[0].finish_reason = this._mapStopReason(event.delta.stop_reason);
|
||||
}
|
||||
if (event.usage) {
|
||||
baseChunk.usage = this._convertUsage(event.usage);
|
||||
}
|
||||
} else if (event.type === 'message_stop') {
|
||||
baseChunk.choices[0].finish_reason = 'stop';
|
||||
}
|
||||
|
||||
return baseChunk;
|
||||
}
|
||||
|
||||
/**
|
||||
* 生成随机 ID
|
||||
*/
|
||||
_generateId() {
|
||||
return Math.random().toString(36).substring(2, 15) +
|
||||
Math.random().toString(36).substring(2, 15);
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = new OpenAIToClaudeConverter();
|
||||
147
src/services/tokenRefreshService.js
Normal file
147
src/services/tokenRefreshService.js
Normal file
@@ -0,0 +1,147 @@
|
||||
const redis = require('../models/redis');
|
||||
const logger = require('../utils/logger');
|
||||
const { v4: uuidv4 } = require('uuid');
|
||||
const {
|
||||
logRefreshSkipped
|
||||
} = require('../utils/tokenRefreshLogger');
|
||||
|
||||
/**
|
||||
* Token 刷新锁服务
|
||||
* 提供分布式锁机制，避免并发刷新问题
|
||||
*/
|
||||
class TokenRefreshService {
|
||||
constructor() {
|
||||
this.lockTTL = 60; // 锁的TTL: 60秒（token刷新通常在30秒内完成）
|
||||
this.lockValue = new Map(); // 存储每个锁的唯一值
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 获取分布式锁
|
||||
* 使用唯一标识符作为值，避免误释放其他进程的锁
|
||||
*/
|
||||
async acquireLock(lockKey) {
|
||||
try {
|
||||
const client = redis.getClientSafe();
|
||||
const lockId = uuidv4();
|
||||
const result = await client.set(lockKey, lockId, 'NX', 'EX', this.lockTTL);
|
||||
|
||||
if (result === 'OK') {
|
||||
this.lockValue.set(lockKey, lockId);
|
||||
logger.debug(`🔒 Acquired lock ${lockKey} with ID ${lockId}, TTL: ${this.lockTTL}s`);
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
} catch (error) {
|
||||
logger.error(`Failed to acquire lock ${lockKey}:`, error);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 释放分布式锁
|
||||
* 使用 Lua 脚本确保只释放自己持有的锁
|
||||
*/
|
||||
async releaseLock(lockKey) {
|
||||
try {
|
||||
const client = redis.getClientSafe();
|
||||
const lockId = this.lockValue.get(lockKey);
|
||||
|
||||
if (!lockId) {
|
||||
logger.warn(`⚠️ No lock ID found for ${lockKey}, skipping release`);
|
||||
return;
|
||||
}
|
||||
|
||||
// Lua 脚本：只有当值匹配时才删除
|
||||
const luaScript = `
|
||||
if redis.call("get", KEYS[1]) == ARGV[1] then
|
||||
return redis.call("del", KEYS[1])
|
||||
else
|
||||
return 0
|
||||
end
|
||||
`;
|
||||
|
||||
const result = await client.eval(luaScript, 1, lockKey, lockId);
|
||||
|
||||
if (result === 1) {
|
||||
this.lockValue.delete(lockKey);
|
||||
logger.debug(`🔓 Released lock ${lockKey} with ID ${lockId}`);
|
||||
} else {
|
||||
logger.warn(`⚠️ Lock ${lockKey} was not released - value mismatch or already expired`);
|
||||
}
|
||||
} catch (error) {
|
||||
logger.error(`Failed to release lock ${lockKey}:`, error);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 获取刷新锁
|
||||
* @param {string} accountId - 账户ID
|
||||
* @param {string} platform - 平台类型 (claude/gemini)
|
||||
* @returns {Promise<boolean>} 是否成功获取锁
|
||||
*/
|
||||
async acquireRefreshLock(accountId, platform = 'claude') {
|
||||
const lockKey = `token_refresh_lock:${platform}:${accountId}`;
|
||||
return await this.acquireLock(lockKey);
|
||||
}
|
||||
|
||||
/**
|
||||
* 释放刷新锁
|
||||
* @param {string} accountId - 账户ID
|
||||
* @param {string} platform - 平台类型 (claude/gemini)
|
||||
*/
|
||||
async releaseRefreshLock(accountId, platform = 'claude') {
|
||||
const lockKey = `token_refresh_lock:${platform}:${accountId}`;
|
||||
await this.releaseLock(lockKey);
|
||||
}
|
||||
|
||||
/**
|
||||
* 检查刷新锁状态
|
||||
* @param {string} accountId - 账户ID
|
||||
* @param {string} platform - 平台类型 (claude/gemini)
|
||||
* @returns {Promise<boolean>} 锁是否存在
|
||||
*/
|
||||
async isRefreshLocked(accountId, platform = 'claude') {
|
||||
const lockKey = `token_refresh_lock:${platform}:${accountId}`;
|
||||
try {
|
||||
const client = redis.getClientSafe();
|
||||
const exists = await client.exists(lockKey);
|
||||
return exists === 1;
|
||||
} catch (error) {
|
||||
logger.error(`Failed to check lock status ${lockKey}:`, error);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 获取锁的剩余TTL
|
||||
* @param {string} accountId - 账户ID
|
||||
* @param {string} platform - 平台类型 (claude/gemini)
|
||||
* @returns {Promise<number>} 剩余秒数，-1表示锁不存在
|
||||
*/
|
||||
async getLockTTL(accountId, platform = 'claude') {
|
||||
const lockKey = `token_refresh_lock:${platform}:${accountId}`;
|
||||
try {
|
||||
const client = redis.getClientSafe();
|
||||
const ttl = await client.ttl(lockKey);
|
||||
return ttl;
|
||||
} catch (error) {
|
||||
logger.error(`Failed to get lock TTL ${lockKey}:`, error);
|
||||
return -1;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 清理本地锁记录
|
||||
* 在进程退出时调用，避免内存泄漏
|
||||
*/
|
||||
cleanup() {
|
||||
this.lockValue.clear();
|
||||
logger.info('🧹 Cleaned up local lock records');
|
||||
}
|
||||
}
|
||||
|
||||
// 创建单例实例
|
||||
const tokenRefreshService = new TokenRefreshService();
|
||||
|
||||
module.exports = tokenRefreshService;
|
||||
Reference in New Issue
Block a user