diff --git a/config/config.go b/config/config.go index bf11b52..de186b0 100644 --- a/config/config.go +++ b/config/config.go @@ -55,6 +55,10 @@ type Account struct { // Priority weight for load balancing (higher = more requests) Weight int `json:"weight,omitempty"` // 0 or 1 = normal, 2+ = higher priority + // Overage behavior after the main usage limit is reached. + AllowOverage bool `json:"allowOverage,omitempty"` // Whether to keep using the account after UsageLimit is reached + OverageWeight int `json:"overageWeight,omitempty"` // 1-10, lower values reduce overage request frequency + // Account status Enabled bool `json:"enabled"` // Whether account is active in the pool BanStatus string `json:"banStatus,omitempty"` // Ban status: "ACTIVE", "BANNED", "SUSPENDED" diff --git a/pool/account.go b/pool/account.go index 2f0ab4d..bba2fad 100644 --- a/pool/account.go +++ b/pool/account.go @@ -9,10 +9,13 @@ import ( "time" ) +const overageFrequencyScale = 10 + // AccountPool 账号池 type AccountPool struct { mu sync.RWMutex accounts []config.Account + totalAccounts int currentIndex uint64 cooldowns map[string]time.Time // 账号冷却时间 errorCounts map[string]int // 连续错误计数 @@ -43,15 +46,19 @@ func (p *AccountPool) Reload() { enabled := config.GetEnabledAccounts() var weighted []config.Account for _, a := range enabled { - w := a.Weight - if w < 1 { - w = 1 + w := effectiveWeight(a.Weight) * overageFrequencyScale + if isOverUsageLimit(a) { + if !a.AllowOverage { + continue + } + w = effectiveOverageWeight(a.OverageWeight) } for j := 0; j < w; j++ { weighted = append(weighted, a) } } p.accounts = weighted + p.totalAccounts = len(enabled) } // GetNext 获取下一个可用账号（加权轮询） @@ -89,7 +96,7 @@ func (p *AccountPool) GetNext() *config.Account { } // 跳过额度已用尽的账号（适用于所有订阅类型） - if acc.UsageLimit > 0 && acc.UsageCurrent >= acc.UsageLimit { + if isOverUsageLimit(*acc) && !acc.AllowOverage { seen[acc.ID] = true continue } @@ -103,7 +110,7 @@ func (p *AccountPool) GetNext() *config.Account { for i := range p.accounts { acc := &p.accounts[i] // 额度用尽的账号不作为 fallback - if acc.UsageLimit > 0 && acc.UsageCurrent >= acc.UsageLimit { + if isOverUsageLimit(*acc) && !acc.AllowOverage { continue } if cooldown, ok := p.cooldowns[acc.ID]; ok { @@ -165,7 +172,6 @@ func (p *AccountPool) UpdateToken(id, accessToken, refreshToken string, expiresA p.accounts[i].RefreshToken = refreshToken } p.accounts[i].ExpiresAt = expiresAt - break } } } @@ -174,7 +180,15 @@ func (p *AccountPool) UpdateToken(id, accessToken, refreshToken string, expiresA func (p *AccountPool) Count() int { p.mu.RLock() defer p.mu.RUnlock() - return len(p.accounts) + if p.totalAccounts > 0 { + return p.totalAccounts + } + + seen := make(map[string]bool) + for _, acc := range p.accounts { + seen[acc.ID] = true + } + return len(seen) } // AvailableCount 返回可用账号数 @@ -183,7 +197,12 @@ func (p *AccountPool) AvailableCount() int { defer p.mu.RUnlock() now := time.Now() count := 0 + seen := make(map[string]bool) for _, acc := range p.accounts { + if seen[acc.ID] { + continue + } + seen[acc.ID] = true if cooldown, ok := p.cooldowns[acc.ID]; ok && now.Before(cooldown) { continue } @@ -196,16 +215,36 @@ func (p *AccountPool) AvailableCount() int { func (p *AccountPool) UpdateStats(id string, tokens int, credits float64) { p.mu.Lock() defer p.mu.Unlock() + var updated bool + var requestCount, errorCount, totalTokens int + var totalCredits float64 + var lastUsed int64 for i := range p.accounts { if p.accounts[i].ID == id { - p.accounts[i].RequestCount++ - p.accounts[i].TotalTokens += tokens - p.accounts[i].TotalCredits += credits - p.accounts[i].LastUsed = time.Now().Unix() - go config.UpdateAccountStats(id, p.accounts[i].RequestCount, p.accounts[i].ErrorCount, p.accounts[i].TotalTokens, p.accounts[i].TotalCredits, p.accounts[i].LastUsed) - break + if !updated { + p.accounts[i].RequestCount++ + p.accounts[i].TotalTokens += tokens + p.accounts[i].TotalCredits += credits + p.accounts[i].LastUsed = time.Now().Unix() + + requestCount = p.accounts[i].RequestCount + errorCount = p.accounts[i].ErrorCount + totalTokens = p.accounts[i].TotalTokens + totalCredits = p.accounts[i].TotalCredits + lastUsed = p.accounts[i].LastUsed + updated = true + continue + } + p.accounts[i].RequestCount = requestCount + p.accounts[i].ErrorCount = errorCount + p.accounts[i].TotalTokens = totalTokens + p.accounts[i].TotalCredits = totalCredits + p.accounts[i].LastUsed = lastUsed } } + if updated { + go config.UpdateAccountStats(id, requestCount, errorCount, totalTokens, totalCredits, lastUsed) + } } // GetAllAccounts 获取所有账号副本 @@ -216,3 +255,24 @@ func (p *AccountPool) GetAllAccounts() []config.Account { copy(result, p.accounts) return result } + +func isOverUsageLimit(acc config.Account) bool { + return acc.UsageLimit > 0 && acc.UsageCurrent >= acc.UsageLimit +} + +func effectiveWeight(weight int) int { + if weight < 1 { + return 1 + } + return weight +} + +func effectiveOverageWeight(weight int) int { + if weight < 1 { + return 1 + } + if weight > overageFrequencyScale { + return overageFrequencyScale + } + return weight +} diff --git a/pool/account_test.go b/pool/account_test.go new file mode 100644 index 0000000..82f5513 --- /dev/null +++ b/pool/account_test.go @@ -0,0 +1,54 @@ +package pool + +import ( + "kiro-api-proxy/config" + "testing" +) + +func TestOverageAccountsAreSkippedByDefault(t *testing.T) { + p := &AccountPool{} + normal := config.Account{ID: "normal"} + overLimit := config.Account{ID: "over", UsageCurrent: 10, UsageLimit: 10} + + p.accounts = []config.Account{normal, overLimit} + + for i := 0; i < 5; i++ { + acc := p.GetNext() + if acc == nil { + t.Fatalf("expected an account") + } + if acc.ID == "over" { + t.Fatalf("expected over-limit account to be skipped by default") + } + } +} + +func TestOverageAccountsCanBeSelectedWhenAllowed(t *testing.T) { + p := &AccountPool{} + overLimit := config.Account{ + ID: "over", + UsageCurrent: 10, + UsageLimit: 10, + AllowOverage: true, + OverageWeight: 1, + } + + p.accounts = []config.Account{overLimit} + + acc := p.GetNext() + if acc == nil { + t.Fatalf("expected allowed overage account") + } + if acc.ID != "over" { + t.Fatalf("expected overage account, got %q", acc.ID) + } +} + +func TestOverageWeightIsLowerThanNormalWeight(t *testing.T) { + normalWeight := effectiveWeight(1) * overageFrequencyScale + overageWeight := effectiveOverageWeight(1) + + if overageWeight >= normalWeight { + t.Fatalf("expected overage weight %d to be lower than normal weight %d", overageWeight, normalWeight) + } +} diff --git a/proxy/handler.go b/proxy/handler.go index aeee7b6..cbfbbcd 100644 --- a/proxy/handler.go +++ b/proxy/handler.go @@ -1969,6 +1969,8 @@ func (h *Handler) apiGetAccounts(w http.ResponseWriter, r *http.Request) { "hasToken": a.AccessToken != "", "machineId": a.MachineId, "weight": a.Weight, + "allowOverage": a.AllowOverage, + "overageWeight": a.OverageWeight, "subscriptionType": a.SubscriptionType, "subscriptionTitle": a.SubscriptionTitle, "daysRemaining": a.DaysRemaining, @@ -2063,6 +2065,12 @@ func (h *Handler) apiUpdateAccount(w http.ResponseWriter, r *http.Request, id st if v, ok := updates["weight"].(float64); ok { existing.Weight = int(v) } + if v, ok := updates["allowOverage"].(bool); ok { + existing.AllowOverage = v + } + if v, ok := updates["overageWeight"].(float64); ok { + existing.OverageWeight = clampInt(int(v), 1, 10) + } if err := config.UpdateAccount(id, *existing); err != nil { w.WriteHeader(500) @@ -2753,6 +2761,9 @@ func (h *Handler) apiGetAccountFull(w http.ResponseWriter, r *http.Request, id s "region": account.Region, "expiresAt": account.ExpiresAt, "machineId": account.MachineId, + "weight": account.Weight, + "allowOverage": account.AllowOverage, + "overageWeight": account.OverageWeight, "enabled": account.Enabled, "banStatus": account.BanStatus, "banReason": account.BanReason, @@ -3111,3 +3122,13 @@ func (h *Handler) apiExportAccounts(w http.ResponseWriter, r *http.Request) { json.NewEncoder(w).Encode(data) } + +func clampInt(v, min, max int) int { + if v < min { + return min + } + if v > max { + return max + } + return v +} diff --git a/web/index.html b/web/index.html index 0caa7e8..c76667e 100644 --- a/web/index.html +++ b/web/index.html @@ -1152,6 +1152,7 @@ 'accounts.copyJSON': '复制 JSON', 'accounts.copyJSONSuccess': 'JSON 已复制到剪贴板', 'accounts.trialDays': '天后到期', + 'accounts.overage': '超额调用', 'time.expired': '已过期', 'time.minutes': '分钟', 'time.hours': '小时', @@ -1344,7 +1345,10 @@ 'filter.banned': '已封禁', 'accounts.weight': '权重', 'detail.weight': '请求权重', - 'detail.weightHint': '0-1=普通, 2+=高优先级' + 'detail.weightHint': '0-1=普通, 2+=高优先级', + 'detail.overage': '超额调用设置', + 'detail.allowOverage': '配额耗尽后继续调用', + 'detail.overageHint': '超额后调用频率权重（1~10），值越大频率越高' }, en: { 'login.subtitle': 'Enter admin password to login', @@ -1384,6 +1388,7 @@ 'accounts.confirmDelete': 'Confirm delete?', 'accounts.copyJSON': 'Copy JSON', 'accounts.copyJSONSuccess': 'JSON copied to clipboard', + 'accounts.overage': 'Overage', 'time.expired': 'Expired', 'time.minutes': 'min', 'time.hours': 'hr', @@ -1572,7 +1577,10 @@ 'filter.banned': 'Banned', 'accounts.weight': 'Weight', 'detail.weight': 'Request Weight', - 'detail.weightHint': '0-1=normal, 2+=higher priority' + 'detail.weightHint': '0-1=normal, 2+=higher priority', + 'detail.overage': 'Overage Settings', + 'detail.allowOverage': 'Continue calling after quota is exhausted', + 'detail.overageHint': 'Call frequency weight when over quota (1–10); higher = more frequent' } }; let currentLang = localStorage.getItem('kiro_lang') || 'zh'; @@ -1838,6 +1846,7 @@ const isSelected = selectedAccounts.has(a.id); const weightVal = a.weight || 0; const weightBadge = weightVal >= 2 ? 'W:' + weightVal + '' : ''; + const overageBadge = a.allowOverage ? '' + t('accounts.overage') + ':' + (a.overageWeight || 1) + '' : ''; return '
' + '
' + '' + @@ -1984,6 +1994,12 @@ '' + t('detail.weightHint') + '' + '' + '
' + + '

' + t('detail.overage') + '

' + + '' + + '' + + '' + t('detail.overageHint') + '' + + '' + + '
' + '

' + t('detail.subscription') + '

' + '
' + t('detail.subscriptionType') + '
' + (a.subscriptionTitle || a.subscriptionType || '-') + '
' + '
' + t('detail.tokenExpiry') + '
' + (a.expiresAt ? new Date(a.expiresAt * 1000).toLocaleString() : '-') + '
' + @@ -2062,6 +2078,20 @@ if (d.success) { alert(t('detail.saved')); loadAccounts(); } else { alert(t('detail.saveFailed') + ': ' + d.error); } } catch (e) { alert(t('detail.saveFailed')); } } + async function saveOverageSettings(id) { + const allowOverage = document.getElementById('allowOverageInput').checked; + let overageWeight = parseInt(document.getElementById('overageWeightInput').value) || 1; + overageWeight = Math.max(1, Math.min(10, overageWeight)); + document.getElementById('overageWeightInput').value = overageWeight; + try { + const res = await fetch('/admin/api/accounts/' + id, { + method: 'PUT', headers: { 'Content-Type': 'application/json', 'X-Admin-Password': password }, + body: JSON.stringify({ allowOverage, overageWeight }) + }); + const d = await res.json(); + if (d.success) { alert(t('detail.saved')); loadAccounts(); } else { alert(t('detail.saveFailed') + ': ' + d.error); } + } catch (e) { alert(t('detail.saveFailed')); } + } async function quickSetWeight(id, value) { const weight = parseInt(value) || 0; try { @@ -2685,4 +2715,4 @@ - \ No newline at end of file +