| const geminiAccountService = require('./geminiAccountService') |
| const accountGroupService = require('./accountGroupService') |
| const redis = require('../models/redis') |
| const logger = require('../utils/logger') |
|
|
| class UnifiedGeminiScheduler { |
| constructor() { |
| this.SESSION_MAPPING_PREFIX = 'unified_gemini_session_mapping:' |
| } |
|
|
| |
| _isSchedulable(schedulable) { |
| |
| if (schedulable === undefined || schedulable === null) { |
| return true |
| } |
| |
| return schedulable !== false && schedulable !== 'false' |
| } |
|
|
| |
| async selectAccountForApiKey(apiKeyData, sessionHash = null, requestedModel = null) { |
| try { |
| |
| if (apiKeyData.geminiAccountId) { |
| |
| if (apiKeyData.geminiAccountId.startsWith('group:')) { |
| const groupId = apiKeyData.geminiAccountId.replace('group:', '') |
| logger.info( |
| `🎯 API key ${apiKeyData.name} is bound to group ${groupId}, selecting from group` |
| ) |
| return await this.selectAccountFromGroup(groupId, sessionHash, requestedModel, apiKeyData) |
| } |
|
|
| |
| const boundAccount = await geminiAccountService.getAccount(apiKeyData.geminiAccountId) |
| if (boundAccount && boundAccount.isActive === 'true' && boundAccount.status !== 'error') { |
| logger.info( |
| `🎯 Using bound dedicated Gemini account: ${boundAccount.name} (${apiKeyData.geminiAccountId}) for API key ${apiKeyData.name}` |
| ) |
| |
| await geminiAccountService.markAccountUsed(apiKeyData.geminiAccountId) |
| return { |
| accountId: apiKeyData.geminiAccountId, |
| accountType: 'gemini' |
| } |
| } else { |
| logger.warn( |
| `⚠️ Bound Gemini account ${apiKeyData.geminiAccountId} is not available, falling back to pool` |
| ) |
| } |
| } |
|
|
| |
| if (sessionHash) { |
| const mappedAccount = await this._getSessionMapping(sessionHash) |
| if (mappedAccount) { |
| |
| const isAvailable = await this._isAccountAvailable( |
| mappedAccount.accountId, |
| mappedAccount.accountType |
| ) |
| if (isAvailable) { |
| |
| await this._extendSessionMappingTTL(sessionHash) |
| logger.info( |
| `🎯 Using sticky session account: ${mappedAccount.accountId} (${mappedAccount.accountType}) for session ${sessionHash}` |
| ) |
| |
| await geminiAccountService.markAccountUsed(mappedAccount.accountId) |
| return mappedAccount |
| } else { |
| logger.warn( |
| `⚠️ Mapped account ${mappedAccount.accountId} is no longer available, selecting new account` |
| ) |
| await this._deleteSessionMapping(sessionHash) |
| } |
| } |
| } |
|
|
| |
| const availableAccounts = await this._getAllAvailableAccounts(apiKeyData, requestedModel) |
|
|
| if (availableAccounts.length === 0) { |
| |
| if (requestedModel) { |
| throw new Error( |
| `No available Gemini accounts support the requested model: ${requestedModel}` |
| ) |
| } else { |
| throw new Error('No available Gemini accounts') |
| } |
| } |
|
|
| |
| const sortedAccounts = this._sortAccountsByPriority(availableAccounts) |
|
|
| |
| const selectedAccount = sortedAccounts[0] |
|
|
| |
| if (sessionHash) { |
| await this._setSessionMapping( |
| sessionHash, |
| selectedAccount.accountId, |
| selectedAccount.accountType |
| ) |
| logger.info( |
| `🎯 Created new sticky session mapping: ${selectedAccount.name} (${selectedAccount.accountId}, ${selectedAccount.accountType}) for session ${sessionHash}` |
| ) |
| } |
|
|
| logger.info( |
| `🎯 Selected account: ${selectedAccount.name} (${selectedAccount.accountId}, ${selectedAccount.accountType}) with priority ${selectedAccount.priority} for API key ${apiKeyData.name}` |
| ) |
|
|
| |
| await geminiAccountService.markAccountUsed(selectedAccount.accountId) |
|
|
| return { |
| accountId: selectedAccount.accountId, |
| accountType: selectedAccount.accountType |
| } |
| } catch (error) { |
| logger.error('❌ Failed to select account for API key:', error) |
| throw error |
| } |
| } |
|
|
| |
| async _getAllAvailableAccounts(apiKeyData, requestedModel = null) { |
| const availableAccounts = [] |
|
|
| |
| if (apiKeyData.geminiAccountId) { |
| const boundAccount = await geminiAccountService.getAccount(apiKeyData.geminiAccountId) |
| if (boundAccount && boundAccount.isActive === 'true' && boundAccount.status !== 'error') { |
| const isRateLimited = await this.isAccountRateLimited(boundAccount.id) |
| if (!isRateLimited) { |
| |
| if ( |
| requestedModel && |
| boundAccount.supportedModels && |
| boundAccount.supportedModels.length > 0 |
| ) { |
| |
| const normalizedModel = requestedModel.replace('models/', '') |
| const modelSupported = boundAccount.supportedModels.some( |
| (model) => model.replace('models/', '') === normalizedModel |
| ) |
| if (!modelSupported) { |
| logger.warn( |
| `⚠️ Bound Gemini account ${boundAccount.name} does not support model ${requestedModel}` |
| ) |
| return availableAccounts |
| } |
| } |
|
|
| logger.info( |
| `🎯 Using bound dedicated Gemini account: ${boundAccount.name} (${apiKeyData.geminiAccountId})` |
| ) |
| return [ |
| { |
| ...boundAccount, |
| accountId: boundAccount.id, |
| accountType: 'gemini', |
| priority: parseInt(boundAccount.priority) || 50, |
| lastUsedAt: boundAccount.lastUsedAt || '0' |
| } |
| ] |
| } |
| } else { |
| logger.warn(`⚠️ Bound Gemini account ${apiKeyData.geminiAccountId} is not available`) |
| } |
| } |
|
|
| |
| const geminiAccounts = await geminiAccountService.getAllAccounts() |
| for (const account of geminiAccounts) { |
| if ( |
| account.isActive === 'true' && |
| account.status !== 'error' && |
| (account.accountType === 'shared' || !account.accountType) && |
| this._isSchedulable(account.schedulable) |
| ) { |
| |
|
|
| |
| const isExpired = geminiAccountService.isTokenExpired(account) |
| if (isExpired && !account.refreshToken) { |
| logger.warn( |
| `⚠️ Gemini account ${account.name} token expired and no refresh token available` |
| ) |
| continue |
| } |
|
|
| |
| if (requestedModel && account.supportedModels && account.supportedModels.length > 0) { |
| |
| const normalizedModel = requestedModel.replace('models/', '') |
| const modelSupported = account.supportedModels.some( |
| (model) => model.replace('models/', '') === normalizedModel |
| ) |
| if (!modelSupported) { |
| logger.debug( |
| `⏭️ Skipping Gemini account ${account.name} - doesn't support model ${requestedModel}` |
| ) |
| continue |
| } |
| } |
|
|
| |
| const isRateLimited = await this.isAccountRateLimited(account.id) |
| if (!isRateLimited) { |
| availableAccounts.push({ |
| ...account, |
| accountId: account.id, |
| accountType: 'gemini', |
| priority: parseInt(account.priority) || 50, |
| lastUsedAt: account.lastUsedAt || '0' |
| }) |
| } |
| } |
| } |
|
|
| logger.info(`📊 Total available Gemini accounts: ${availableAccounts.length}`) |
| return availableAccounts |
| } |
|
|
| |
| _sortAccountsByPriority(accounts) { |
| return accounts.sort((a, b) => { |
| |
| if (a.priority !== b.priority) { |
| return a.priority - b.priority |
| } |
|
|
| |
| const aLastUsed = new Date(a.lastUsedAt || 0).getTime() |
| const bLastUsed = new Date(b.lastUsedAt || 0).getTime() |
| return aLastUsed - bLastUsed |
| }) |
| } |
|
|
| |
| async _isAccountAvailable(accountId, accountType) { |
| try { |
| if (accountType === 'gemini') { |
| const account = await geminiAccountService.getAccount(accountId) |
| if (!account || account.isActive !== 'true' || account.status === 'error') { |
| return false |
| } |
| |
| if (!this._isSchedulable(account.schedulable)) { |
| logger.info(`🚫 Gemini account ${accountId} is not schedulable`) |
| return false |
| } |
| return !(await this.isAccountRateLimited(accountId)) |
| } |
| return false |
| } catch (error) { |
| logger.warn(`⚠️ Failed to check account availability: ${accountId}`, error) |
| return false |
| } |
| } |
|
|
| |
| async _getSessionMapping(sessionHash) { |
| const client = redis.getClientSafe() |
| const mappingData = await client.get(`${this.SESSION_MAPPING_PREFIX}${sessionHash}`) |
|
|
| if (mappingData) { |
| try { |
| return JSON.parse(mappingData) |
| } catch (error) { |
| logger.warn('⚠️ Failed to parse session mapping:', error) |
| return null |
| } |
| } |
|
|
| return null |
| } |
|
|
| |
| async _setSessionMapping(sessionHash, accountId, accountType) { |
| const client = redis.getClientSafe() |
| const mappingData = JSON.stringify({ accountId, accountType }) |
| |
| const appConfig = require('../../config/config') |
| const ttlHours = appConfig.session?.stickyTtlHours || 1 |
| const ttlSeconds = Math.max(1, Math.floor(ttlHours * 60 * 60)) |
| await client.setex(`${this.SESSION_MAPPING_PREFIX}${sessionHash}`, ttlSeconds, mappingData) |
| } |
|
|
| |
| async _deleteSessionMapping(sessionHash) { |
| const client = redis.getClientSafe() |
| await client.del(`${this.SESSION_MAPPING_PREFIX}${sessionHash}`) |
| } |
|
|
| |
| async _extendSessionMappingTTL(sessionHash) { |
| try { |
| const client = redis.getClientSafe() |
| const key = `${this.SESSION_MAPPING_PREFIX}${sessionHash}` |
| const remainingTTL = await client.ttl(key) |
|
|
| if (remainingTTL === -2) { |
| return false |
| } |
| if (remainingTTL === -1) { |
| return true |
| } |
|
|
| const appConfig = require('../../config/config') |
| const ttlHours = appConfig.session?.stickyTtlHours || 1 |
| const renewalThresholdMinutes = appConfig.session?.renewalThresholdMinutes || 0 |
| if (!renewalThresholdMinutes) { |
| return true |
| } |
|
|
| const fullTTL = Math.max(1, Math.floor(ttlHours * 60 * 60)) |
| const threshold = Math.max(0, Math.floor(renewalThresholdMinutes * 60)) |
|
|
| if (remainingTTL < threshold) { |
| await client.expire(key, fullTTL) |
| logger.debug( |
| `🔄 Renewed unified Gemini session TTL: ${sessionHash} (was ${Math.round(remainingTTL / 60)}m, renewed to ${ttlHours}h)` |
| ) |
| } else { |
| logger.debug( |
| `✅ Unified Gemini session TTL sufficient: ${sessionHash} (remaining ${Math.round(remainingTTL / 60)}m)` |
| ) |
| } |
| return true |
| } catch (error) { |
| logger.error('❌ Failed to extend unified Gemini session TTL:', error) |
| return false |
| } |
| } |
|
|
| |
| async markAccountRateLimited(accountId, accountType, sessionHash = null) { |
| try { |
| if (accountType === 'gemini') { |
| await geminiAccountService.setAccountRateLimited(accountId, true) |
| } |
|
|
| |
| if (sessionHash) { |
| await this._deleteSessionMapping(sessionHash) |
| } |
|
|
| return { success: true } |
| } catch (error) { |
| logger.error( |
| `❌ Failed to mark account as rate limited: ${accountId} (${accountType})`, |
| error |
| ) |
| throw error |
| } |
| } |
|
|
| |
| async removeAccountRateLimit(accountId, accountType) { |
| try { |
| if (accountType === 'gemini') { |
| await geminiAccountService.setAccountRateLimited(accountId, false) |
| } |
|
|
| return { success: true } |
| } catch (error) { |
| logger.error( |
| `❌ Failed to remove rate limit for account: ${accountId} (${accountType})`, |
| error |
| ) |
| throw error |
| } |
| } |
|
|
| |
| async isAccountRateLimited(accountId) { |
| try { |
| const account = await geminiAccountService.getAccount(accountId) |
| if (!account) { |
| return false |
| } |
|
|
| if (account.rateLimitStatus === 'limited' && account.rateLimitedAt) { |
| const limitedAt = new Date(account.rateLimitedAt).getTime() |
| const now = Date.now() |
| const limitDuration = 60 * 60 * 1000 |
|
|
| return now < limitedAt + limitDuration |
| } |
| return false |
| } catch (error) { |
| logger.error(`❌ Failed to check rate limit status: ${accountId}`, error) |
| return false |
| } |
| } |
|
|
| |
| async selectAccountFromGroup(groupId, sessionHash = null, requestedModel = null) { |
| try { |
| |
| const group = await accountGroupService.getGroup(groupId) |
| if (!group) { |
| throw new Error(`Group ${groupId} not found`) |
| } |
|
|
| if (group.platform !== 'gemini') { |
| throw new Error(`Group ${group.name} is not a Gemini group`) |
| } |
|
|
| logger.info(`👥 Selecting account from Gemini group: ${group.name}`) |
|
|
| |
| if (sessionHash) { |
| const mappedAccount = await this._getSessionMapping(sessionHash) |
| if (mappedAccount) { |
| |
| const memberIds = await accountGroupService.getGroupMembers(groupId) |
| if (memberIds.includes(mappedAccount.accountId)) { |
| const isAvailable = await this._isAccountAvailable( |
| mappedAccount.accountId, |
| mappedAccount.accountType |
| ) |
| if (isAvailable) { |
| |
| await this._extendSessionMappingTTL(sessionHash) |
| logger.info( |
| `🎯 Using sticky session account from group: ${mappedAccount.accountId} (${mappedAccount.accountType}) for session ${sessionHash}` |
| ) |
| |
| await geminiAccountService.markAccountUsed(mappedAccount.accountId) |
| return mappedAccount |
| } |
| } |
| |
| await this._deleteSessionMapping(sessionHash) |
| } |
| } |
|
|
| |
| const memberIds = await accountGroupService.getGroupMembers(groupId) |
| if (memberIds.length === 0) { |
| throw new Error(`Group ${group.name} has no members`) |
| } |
|
|
| const availableAccounts = [] |
|
|
| |
| for (const memberId of memberIds) { |
| const account = await geminiAccountService.getAccount(memberId) |
|
|
| if (!account) { |
| logger.warn(`⚠️ Gemini account ${memberId} not found in group ${group.name}`) |
| continue |
| } |
|
|
| |
| if ( |
| account.isActive === 'true' && |
| account.status !== 'error' && |
| this._isSchedulable(account.schedulable) |
| ) { |
| |
| const isExpired = geminiAccountService.isTokenExpired(account) |
| if (isExpired && !account.refreshToken) { |
| logger.warn( |
| `⚠️ Gemini account ${account.name} in group token expired and no refresh token available` |
| ) |
| continue |
| } |
|
|
| |
| if (requestedModel && account.supportedModels && account.supportedModels.length > 0) { |
| |
| const normalizedModel = requestedModel.replace('models/', '') |
| const modelSupported = account.supportedModels.some( |
| (model) => model.replace('models/', '') === normalizedModel |
| ) |
| if (!modelSupported) { |
| logger.debug( |
| `⏭️ Skipping Gemini account ${account.name} in group - doesn't support model ${requestedModel}` |
| ) |
| continue |
| } |
| } |
|
|
| |
| const isRateLimited = await this.isAccountRateLimited(account.id) |
| if (!isRateLimited) { |
| availableAccounts.push({ |
| ...account, |
| accountId: account.id, |
| accountType: 'gemini', |
| priority: parseInt(account.priority) || 50, |
| lastUsedAt: account.lastUsedAt || '0' |
| }) |
| } |
| } |
| } |
|
|
| if (availableAccounts.length === 0) { |
| throw new Error(`No available accounts in Gemini group ${group.name}`) |
| } |
|
|
| |
| const sortedAccounts = this._sortAccountsByPriority(availableAccounts) |
|
|
| |
| const selectedAccount = sortedAccounts[0] |
|
|
| |
| if (sessionHash) { |
| await this._setSessionMapping( |
| sessionHash, |
| selectedAccount.accountId, |
| selectedAccount.accountType |
| ) |
| logger.info( |
| `🎯 Created new sticky session mapping in group: ${selectedAccount.name} (${selectedAccount.accountId}, ${selectedAccount.accountType}) for session ${sessionHash}` |
| ) |
| } |
|
|
| logger.info( |
| `🎯 Selected account from Gemini group ${group.name}: ${selectedAccount.name} (${selectedAccount.accountId}, ${selectedAccount.accountType}) with priority ${selectedAccount.priority}` |
| ) |
|
|
| |
| await geminiAccountService.markAccountUsed(selectedAccount.accountId) |
|
|
| return { |
| accountId: selectedAccount.accountId, |
| accountType: selectedAccount.accountType |
| } |
| } catch (error) { |
| logger.error(`❌ Failed to select account from Gemini group ${groupId}:`, error) |
| throw error |
| } |
| } |
| } |
|
|
| module.exports = new UnifiedGeminiScheduler() |
|
|