diff --git a/src/i18n/authFileList.ts b/src/i18n/authFileList.ts index 08e45e1a..44fc6be5 100644 --- a/src/i18n/authFileList.ts +++ b/src/i18n/authFileList.ts @@ -1,5 +1,7 @@ export const authFileListZhCN = { 'quota.pluginCategory': '插件额度', + 'quota.autoBalance': '按剩余额度自动均衡权重', + 'quota.autoBalanceHint': '刷新额度时根据剩余比例自动更新凭证的 priority 与 weight', "authFiles.policy.title": "请求策略", "authFiles.policy.hint": "未设置时继承全局配置。错误规则按顺序优先生效,清空后恢复继承;正则使用 Go 语法。", "authFiles.operations.title": "批量管理", @@ -43,6 +45,8 @@ export const authFileListZhCN = { export const authFileListEn: Record = { 'quota.pluginCategory': 'Plugin quotas', + 'quota.autoBalance': 'Auto-balance weights', + 'quota.autoBalanceHint': 'Automatically update credential priority and weight based on remaining quota when refreshed', "authFiles.policy.title": "Request policies", "authFiles.policy.hint": "Unset values inherit global settings. Error rules run in order; clearing restores inheritance. Regex uses Go syntax.", "authFiles.operations.title": "Batch management", @@ -86,6 +90,8 @@ export const authFileListEn: Record = { export const authFileListJa: Record = { 'quota.pluginCategory': 'プラグイン割当量', + 'quota.autoBalance': '残量に応じて重みを自動調整', + 'quota.autoBalanceHint': 'クォータ更新時に残り割合に応じて認証情報の priority と weight を自動更新します', "authFiles.policy.title": "リクエストポリシー", "authFiles.policy.hint": "未設定は全体設定を継承します。ルールは順番に適用され、空にすると継承に戻ります。正規表現は Go 構文です。", "authFiles.operations.title": "一括管理", diff --git a/src/pages/AuthFileManagementPage.tsx b/src/pages/AuthFileManagementPage.tsx index 54950830..51555f34 100644 --- a/src/pages/AuthFileManagementPage.tsx +++ b/src/pages/AuthFileManagementPage.tsx @@ -3,6 +3,7 @@ import { refreshCredential } from '../services/authFileOperations'; import { EmptyState } from '../components/EmptyState'; import { Checkbox } from '../components/Checkbox'; import { commitFetchedQuota } from '../services/quotaEnrichment'; +import { applyQuotaRoutingChange, subscribeQuotaRoutingChanges } from '../services/quotaAutoBalance'; import { ChangeEvent, useCallback, useEffect, useMemo, useRef, useState } from 'react'; import { useConfirmation } from '../components/ConfirmationDialog'; import { QuotaActionFeedback } from '../components/QuotaActionFeedback'; @@ -137,6 +138,12 @@ export function AuthFileManagementPage({ initialProvider }: { initialProvider?: observedAt?: string; }>({ files: [], receivedAtMs: 0 }); const { files, receivedAtMs, observedAt } = fileSnapshot; + useEffect(() => subscribeQuotaRoutingChanges((name, fields) => { + setFileSnapshot((current) => { + const nextFiles = applyQuotaRoutingChange(current.files, name, fields); + return nextFiles === current.files ? current : { ...current, files: nextFiles }; + }); + }), []); const [filter, setFilter] = useState(''); const [page, setPage] = useState(1); const [pageSize, setPageSize] = useState(10); diff --git a/src/pages/QuotaPage.tsx b/src/pages/QuotaPage.tsx index deea425c..891112b3 100644 --- a/src/pages/QuotaPage.tsx +++ b/src/pages/QuotaPage.tsx @@ -1,7 +1,15 @@ import { Notice } from '../components/Notice'; import { EmptyState } from '../components/EmptyState'; import { SegmentedControl } from '../components/SegmentedControl'; +import { Switch } from '../components/Switch'; import { commitFetchedQuota } from '../services/quotaEnrichment'; +import { + applyQuotaRoutingChange, + restoreQuotaRoutingOriginals, + subscribeQuotaRoutingChanges, + syncCredentialQuotaRouting, + useQuotaAutoBalance, +} from '../services/quotaAutoBalance'; import { ResetCreditExpiries } from '../components/ResetCreditExpiries'; import { MessageNotice } from '../appNotice'; import { useCallback, useEffect, useMemo, useState } from 'react'; @@ -64,8 +72,28 @@ export function QuotaPage() { const [refreshing, setRefreshing] = useState(false); const [error, setError] = useState(''); const [category, setCategory] = useState<'all' | QuotaDisplayProvider>('all'); + const [autoBalance, setAutoBalance] = useQuotaAutoBalance(); const querying = Object.values(quotas).some((quota) => quota.status === 'loading'); + useEffect(() => subscribeQuotaRoutingChanges((name, fields) => { + setFiles((current) => applyQuotaRoutingChange(current, name, fields)); + }), []); + + const handleToggleAutoBalance = useCallback((enabled: boolean) => { + setAutoBalance(enabled); + if (!enabled) { + void restoreQuotaRoutingOriginals().catch((restoreError) => setError(String(restoreError))); + return; + } + const snapshot = getQuotaCacheSnapshot(); + files.forEach((file) => { + const quota = snapshot[quotaKey(file)]; + if (quota?.status === 'success') { + void syncCredentialQuotaRouting(file, quota).catch(() => {}); + } + }); + }, [files, setAutoBalance]); + const loadFiles = useCallback(async () => { setLoading(true); setError(''); @@ -156,6 +184,13 @@ export function QuotaPage() { icon={} summary={t(files.length === 1 ? 'quota.queryableCredentials.one' : 'quota.queryableCredentials.other', { count: files.length })} > + handleToggleAutoBalance(event.currentTarget.checked)} + label={t('quota.autoBalance')} + title={t('quota.autoBalanceHint')} + /> diff --git a/src/services/quotaAutoBalance.ts b/src/services/quotaAutoBalance.ts new file mode 100644 index 00000000..a6b53222 --- /dev/null +++ b/src/services/quotaAutoBalance.ts @@ -0,0 +1,348 @@ +import { useSyncExternalStore } from 'react'; +import { isOAuthCredentialFile, normalizeAuthFilePriorityInput } from './authFiles'; +import { isRecord, managementApi, readBoolean, readString, responseList } from './managementApi'; +import type { AuthFile, QuotaRow, QuotaState } from './quotaService'; + +const STORAGE_KEY = 'easy-cli-proxy-api.quota-auto-balance'; +const ORIGINALS_STORAGE_KEY = 'easy-cli-proxy-api.quota-auto-balance.originals'; + +/** + * Priorities are anchored to the backend default (0) so credentials auto-balance never + * touched rank alongside healthy ones, and depleted credentials rank below both. + */ +export const AUTO_BALANCE_PRIORITY = { healthy: 0, low: -5, critical: -10 } as const; + +let cachedEnabled: boolean | undefined; +const listeners = new Set<() => void>(); + +/** Returns undefined when storage is unavailable so the in-memory value stays authoritative. */ +function readStoredEnabled(): boolean | undefined { + try { + return localStorage.getItem(STORAGE_KEY) === 'on'; + } catch { + return undefined; + } +} + +export function getQuotaAutoBalanceEnabled(): boolean { + if (cachedEnabled === undefined) { + cachedEnabled = readStoredEnabled() ?? false; + } + return cachedEnabled; +} + +/** Re-reads the stored toggle, which another tab may have changed, and notifies on change. */ +function refreshQuotaAutoBalanceEnabled(): boolean { + const stored = readStoredEnabled(); + if (stored !== undefined && stored !== getQuotaAutoBalanceEnabled()) { + cachedEnabled = stored; + listeners.forEach((listener) => listener()); + } + return getQuotaAutoBalanceEnabled(); +} + +/** Keeps the toggle in step with other tabs; `key` is null when storage was cleared. */ +export function watchQuotaAutoBalanceStorage(target: Pick): () => void { + const onStorage = (event: Event) => { + const key = (event as Event & { key?: string | null }).key; + if (key === STORAGE_KEY || key === null || key === undefined) refreshQuotaAutoBalanceEnabled(); + }; + target.addEventListener('storage', onStorage); + return () => target.removeEventListener('storage', onStorage); +} + +if (typeof window !== 'undefined' && typeof window.addEventListener === 'function') { + watchQuotaAutoBalanceStorage(window); +} + +export function setQuotaAutoBalanceEnabled(enabled: boolean): void { + const next = Boolean(enabled); + const changed = getQuotaAutoBalanceEnabled() !== next; + cachedEnabled = next; + try { + localStorage.setItem(STORAGE_KEY, next ? 'on' : 'off'); + } catch { + // Ignore storage errors in restricted environments. + } + if (changed) { + listeners.forEach((listener) => listener()); + } +} + +function subscribeQuotaAutoBalance(listener: () => void): () => void { + listeners.add(listener); + return () => { + listeners.delete(listener); + }; +} + +export function useQuotaAutoBalance(): [boolean, (enabled: boolean) => void] { + const enabled = useSyncExternalStore( + subscribeQuotaAutoBalance, + getQuotaAutoBalanceEnabled, + getQuotaAutoBalanceEnabled, + ); + return [enabled, setQuotaAutoBalanceEnabled]; +} + +export type QuotaRoutingFields = { priority: number; weight: number | null }; +type RoutingChangeListener = (name: string, fields: QuotaRoutingFields) => void; +const routingChangeListeners = new Set(); + +/** Notifies pages holding credential lists so they can update rows immutably after a PATCH. */ +export function subscribeQuotaRoutingChanges(listener: RoutingChangeListener): () => void { + routingChangeListeners.add(listener); + return () => { + routingChangeListeners.delete(listener); + }; +} + +export function applyQuotaRoutingChange(files: T[], name: string, fields: QuotaRoutingFields): T[] { + let changed = false; + const next = files.map((file) => { + if (readString(file, 'name') !== name) return file; + changed = true; + return { ...file, priority: fields.priority, weight: fields.weight }; + }); + return changed ? next : files; +} + +/** + * Manual priority and weight captured before auto-balance first changed a credential, plus the + * values auto-balance last applied so a later manual edit can be told apart from its own writes. + */ +type StoredOriginal = QuotaRoutingFields & { applied?: QuotaRoutingFields }; +type StoredOriginals = Record; + +const isRoutingFields = (value: unknown): value is QuotaRoutingFields => isRecord(value) + && typeof value.priority === 'number' + && (value.weight === null || typeof value.weight === 'number'); + +function readOriginals(): StoredOriginals { + try { + const parsed: unknown = JSON.parse(localStorage.getItem(ORIGINALS_STORAGE_KEY) ?? '{}'); + if (!isRecord(parsed)) return {}; + return Object.fromEntries(Object.entries(parsed).filter(([, value]) => isRoutingFields(value) + && ((value as StoredOriginal).applied === undefined || isRoutingFields((value as StoredOriginal).applied)))) as StoredOriginals; + } catch { + return {}; + } +} + +function writeOriginals(originals: StoredOriginals): void { + try { + if (Object.keys(originals).length) localStorage.setItem(ORIGINALS_STORAGE_KEY, JSON.stringify(originals)); + else localStorage.removeItem(ORIGINALS_STORAGE_KEY); + } catch { + // Ignore storage errors in restricted environments. + } +} + +function updateOriginal(name: string, next: StoredOriginal | undefined): void { + const originals = readOriginals(); + if (next) originals[name] = next; + else delete originals[name]; + writeOriginals(originals); +} + +export type QuotaRoutingTarget = { + priority: number; + weight: number; + bottleneckPercent: number; + averagePercent: number; +}; + +const clampedPercent = (row: QuotaRow): number | null => + typeof row.remainingPercent === 'number' && Number.isFinite(row.remainingPercent) + ? Math.max(0, Math.min(100, row.remainingPercent)) + : null; + +/** + * Percentages that gate the whole credential. Account-wide windows (5-hour, weekly, ...) all apply, + * so the tightest one wins. Model- or feature-scoped limits only block part of the traffic, so they + * are ignored when account-wide windows exist; otherwise (e.g. Antigravity model groups) each group + * is reduced to its tightest window and the best group represents the credential. + */ +function routingPercents(rows: QuotaRow[]): { bottleneck: number; values: number[] } | null { + const account: number[] = []; + const groups = new Map(); + for (const row of rows) { + const percent = clampedPercent(row); + if (percent === null) continue; + if (!row.scoped && !row.groupLabel) { + account.push(percent); + continue; + } + const group = row.groupLabel || row.label; + groups.set(group, Math.min(groups.get(group) ?? 100, percent)); + } + if (account.length) return { bottleneck: Math.min(...account), values: account }; + const values = [...groups.values()]; + return values.length ? { bottleneck: Math.max(...values), values } : null; +} + +/** + * Compute recommended priority and weight from a fetched QuotaState. + * Returns null if the quota is not successful or contains no numeric percentage rows. + */ +export function computeQuotaRoutingTarget(quota: QuotaState): QuotaRoutingTarget | null { + if (quota.status !== 'success' || !Array.isArray(quota.rows) || quota.rows.length === 0) { + return null; + } + const percents = routingPercents(quota.rows); + if (!percents) return null; + + const bottleneckPercent = percents.bottleneck; + const averagePercent = percents.values.reduce((acc, value) => acc + value, 0) / percents.values.length; + + let priority: number; + let weight: number; + + if (bottleneckPercent < 5) { + priority = AUTO_BALANCE_PRIORITY.critical; + weight = 1; + } else if (bottleneckPercent < 15) { + priority = AUTO_BALANCE_PRIORITY.low; + weight = Math.max(1, Math.round(bottleneckPercent)); + } else { + priority = AUTO_BALANCE_PRIORITY.healthy; + const effective = bottleneckPercent * 0.7 + averagePercent * 0.3; + weight = Math.max(1, Math.round(effective)); + } + + return { + priority, + weight, + bottleneckPercent, + averagePercent, + }; +} + +type BalanceApi = { + get: (path: string, query?: Record) => Promise; + patch: (path: string, body: Record) => Promise; +}; + +const parseWeight = (value: unknown): number | null => { + const text = typeof value === 'number' || typeof value === 'string' ? String(value).trim() : ''; + if (!text) return null; + const weight = Number(text); + return Number.isFinite(weight) ? Math.max(0, weight) : null; +}; + +const sameRouting = (left: QuotaRoutingFields, right: QuotaRoutingFields) => + left.priority === right.priority && left.weight === right.weight; + +/** Reads only the routing fields, so unrelated malformed metadata never blocks balancing. */ +async function loadRoutingFields(api: BalanceApi, name: string): Promise { + let metadata: unknown = await api.get('/credentials/download', { name }); + if (typeof metadata === 'string') metadata = JSON.parse(metadata); + if (!isRecord(metadata)) throw new Error(`Unexpected credential payload for ${name}`); + const priority = metadata.priority; + return { + priority: typeof priority === 'number' || typeof priority === 'string' + ? normalizeAuthFilePriorityInput(String(priority)) ?? 0 + : 0, + weight: parseWeight(metadata.weight), + }; +} + +async function patchRouting(api: BalanceApi, name: string, fields: QuotaRoutingFields): Promise { + await api.patch('/credentials/fields', { name, ...fields }); + routingChangeListeners.forEach((listener) => listener(name, fields)); +} + +/** Sync and restore of one credential run one at a time so a late sync cannot undo a restore. */ +const credentialLocks = new Map>(); + +function withCredentialLock(name: string, task: () => Promise): Promise { + const previous = credentialLocks.get(name) ?? Promise.resolve(); + const run = previous.catch(() => undefined).then(task); + const tail = run.catch(() => undefined); + credentialLocks.set(name, tail); + void tail.then(() => { + if (credentialLocks.get(name) === tail) credentialLocks.delete(name); + }); + return run; +} + +/** + * Apply computed routing priority and weight to an OAuth credential file if changed. + * The credential list does not carry weight, so current values are read from the stored file. + */ +export async function syncCredentialQuotaRouting( + file: AuthFile, + quota: QuotaState, + api: BalanceApi = managementApi, +): Promise { + if (!isOAuthCredentialFile(file) || readBoolean(file, 'disabled')) { + return null; + } + const name = readString(file, 'name'); + if (!name) return null; + + const target = computeQuotaRoutingTarget(quota); + if (!target) return null; + + return withCredentialLock(name, async () => { + // Read the stored toggle, not just this tab's copy, so another tab turning it off wins. + if (!refreshQuotaAutoBalanceEnabled()) return null; + const current = await loadRoutingFields(api, name); + if (!refreshQuotaAutoBalanceEnabled()) return null; + + const desired = { priority: target.priority, weight: target.weight }; + const stored = readOriginals()[name]; + // A value differing from what auto-balance last applied is a newer manual setting. + const manual = stored?.applied && sameRouting(current, stored.applied) + ? { priority: stored.priority, weight: stored.weight } + : current; + if (sameRouting(current, desired)) { + updateOriginal(name, stored?.applied ? { ...manual, applied: desired } : { ...manual }); + return target; + } + + updateOriginal(name, { ...manual, applied: desired }); + try { + await patchRouting(api, name, desired); + } catch (error) { + updateOriginal(name, stored); + throw error; + } + return target; + }); +} + +/** + * Restore the manual priority and weight of every credential auto-balance changed. Credentials + * that no longer exist are forgotten, and ones edited by hand since the last sync are left alone. + */ +export async function restoreQuotaRoutingOriginals(api: BalanceApi = managementApi): Promise { + const originals = readOriginals(); + if (!Object.keys(originals).length) return; + const existing = new Set(responseList(await api.get('/credentials'), 'files').map((file) => readString(file, 'name'))); + + const failures: unknown[] = []; + for (const name of Object.keys(originals)) { + if (!existing.has(name)) { + updateOriginal(name, undefined); + continue; + } + try { + await withCredentialLock(name, async () => { + if (refreshQuotaAutoBalanceEnabled()) return; + const stored = readOriginals()[name]; + if (!stored) return; + if (stored.applied) { + const current = await loadRoutingFields(api, name); + if (sameRouting(current, stored.applied)) { + await patchRouting(api, name, { priority: stored.priority, weight: stored.weight }); + } + } + updateOriginal(name, undefined); + }); + } catch (error) { + failures.push(error); + } + } + if (failures.length) throw failures[0]; +} diff --git a/src/services/quotaEnrichment.ts b/src/services/quotaEnrichment.ts index 7dff9f05..14936b17 100644 --- a/src/services/quotaEnrichment.ts +++ b/src/services/quotaEnrichment.ts @@ -1,3 +1,4 @@ +import { getQuotaAutoBalanceEnabled, syncCredentialQuotaRouting } from './quotaAutoBalance'; import { enrichXaiQuotaPlan, providerForFile, quotaKey, type AuthFile, type QuotaState } from './quotaService'; import { commitQuotaCacheIfCurrent, getQuotaCacheSnapshot, updateQuotaCache } from './quotaCache'; @@ -6,12 +7,19 @@ export function commitFetchedQuota(generation: number, file: AuthFile, quota: Qu const committed = commitQuotaCacheIfCurrent(generation, () => { updateQuotaCache((current) => ({ ...current, [key]: quota })); }); - if (committed && quota.status === 'success' && providerForFile(file) === 'xai') { - void enrichXaiQuotaPlan(file, quota).then((enriched) => { - if (enriched === quota || getQuotaCacheSnapshot()[key] !== quota) return; - commitQuotaCacheIfCurrent(generation, () => { - updateQuotaCache((current) => ({ ...current, [key]: enriched })); + if (committed && quota.status === 'success') { + if (getQuotaAutoBalanceEnabled()) { + void syncCredentialQuotaRouting(file, quota).catch(() => { + /* Keep quota updates resilient if routing sync fails. */ }); - }).catch(() => { /* Optional plan failures preserve the original quota. */ }); + } + if (providerForFile(file) === 'xai') { + void enrichXaiQuotaPlan(file, quota).then((enriched) => { + if (enriched === quota || getQuotaCacheSnapshot()[key] !== quota) return; + commitQuotaCacheIfCurrent(generation, () => { + updateQuotaCache((current) => ({ ...current, [key]: enriched })); + }); + }).catch(() => { /* Optional plan failures preserve the original quota. */ }); + } } } diff --git a/src/services/quotaService.ts b/src/services/quotaService.ts index 46b13178..43f6ec8c 100644 --- a/src/services/quotaService.ts +++ b/src/services/quotaService.ts @@ -30,6 +30,8 @@ export type QuotaRow = { label: string; remainingPercent: number | null; quotaAvailable?: boolean; + /** Limit applies to a subset of models or features rather than the whole account. */ + scoped?: boolean; groupLabel?: string; groupDescription?: string; windowLabel?: string; @@ -68,7 +70,9 @@ const endpointByProvider: Record = { antigravity: 'https://daily-cloudcode-pa.googleapis.com/v1internal:retrieveUserQuotaSummary', }; -const CLAUDE_PROFILE_URL = 'https://api.anthropic.com/api/oauth/profile'; +/** Claude windows that gate all proxy (OAuth app) traffic rather than specific models or products. */ +const ACCOUNT_WIDE_CLAUDE_WINDOWS = new Set(['five_hour', 'seven_day', 'seven_day_oauth_apps']); +const CLAUDE_PROFILE_URL ='https://api.anthropic.com/api/oauth/profile'; const XAI_WEEKLY_URL = 'https://cli-chat-proxy.grok.com/v1/billing?format=credits'; const CODEX_RESET_CREDITS_URL = 'https://chatgpt.com/backend-api/wham/rate-limit-reset-credits'; @@ -307,6 +311,7 @@ const codexWindowRows = (value: Record): QuotaRow[] => { ?? (reached && resetAtMs !== undefined ? 0 : null), reset: codexResetLabel(raw), resetAtMs, + ...(prefix ? { scoped: true } : {}), }; }).filter((row): row is QuotaRow => row !== null); }; @@ -446,6 +451,7 @@ export const quotaRowsFor = (provider: QuotaProvider, payload: unknown): QuotaRo remainingPercent: remainingFromUsedPercent(raw.utilization), reset: absoluteResetLabel(raw.resets_at ?? raw.resetsAt), resetAtMs: quotaResetFor(raw, ['resets_at', 'resetsAt']), + ...(ACCOUNT_WIDE_CLAUDE_WINDOWS.has(key) ? {} : { scoped: true }), }; }) .filter((row): row is QuotaRow => row !== null); @@ -454,6 +460,7 @@ export const quotaRowsFor = (provider: QuotaProvider, payload: unknown): QuotaRo remainingPercent: remainingFromUsedPercent(fable.percent), reset: absoluteResetLabel(fable.resets_at ?? fable.resetsAt), resetAtMs: quotaResetFor(fable, ['resets_at', 'resetsAt']), + scoped: true, }); const extraUsage = isRecord(value.extra_usage) ? value.extra_usage @@ -476,6 +483,7 @@ export const quotaRowsFor = (provider: QuotaProvider, payload: unknown): QuotaRo detail: usedLabel && limitLabel ? quotaText('quota.service.usedOf', { used: usedLabel, limit: limitLabel }) : undefined, + scoped: true, }); } return rows; diff --git a/tests/quotaAutoBalance.test.ts b/tests/quotaAutoBalance.test.ts new file mode 100644 index 00000000..69426bf2 --- /dev/null +++ b/tests/quotaAutoBalance.test.ts @@ -0,0 +1,306 @@ +import { afterEach, beforeEach, describe, expect, it, spyOn } from 'bun:test'; +import { managementApi } from '../src/services/managementApi'; +import { + AUTO_BALANCE_PRIORITY, + applyQuotaRoutingChange, + computeQuotaRoutingTarget, + getQuotaAutoBalanceEnabled, + restoreQuotaRoutingOriginals, + setQuotaAutoBalanceEnabled, + subscribeQuotaRoutingChanges, + syncCredentialQuotaRouting, + watchQuotaAutoBalanceStorage, +} from '../src/services/quotaAutoBalance'; +import { commitFetchedQuota } from '../src/services/quotaEnrichment'; +import { captureQuotaCacheGeneration, updateQuotaCache } from '../src/services/quotaCache'; +import { quotaRowsFor, type QuotaState } from '../src/services/quotaService'; + +const memoryStorage = () => { + const values = new Map(); + return { + getItem: (key: string) => values.get(key) ?? null, + setItem: (key: string, value: string) => { values.set(key, value); }, + removeItem: (key: string) => { values.delete(key); }, + }; +}; + +const fakeApi = (stored: Record>) => { + const patches: Array> = []; + const api = { + patches, + beforeDownload: undefined as undefined | (() => Promise), + get: async (path: string, query?: Record) => { + if (path === '/credentials') return { files: Object.keys(stored).map((name) => ({ name })) }; + await api.beforeDownload?.(); + return { ...stored[String(query?.name)] }; + }, + patch: async (_path: string, body: Record) => { + patches.push(body); + const { name, ...fields } = body; + Object.assign(stored[String(name)], fields); + }, + }; + return api; +}; + +const claudeFile = { name: 'claude-test.json', provider: 'claude', auth_index: 'idx-1', source: 'file', priority: 3 }; + +describe('quotaAutoBalance', () => { + const originalStorage = (globalThis as { localStorage?: unknown }).localStorage; + let spies: Array<{ mockRestore: () => void }> = []; + + beforeEach(() => { + (globalThis as { localStorage?: unknown }).localStorage = memoryStorage(); + }); + + afterEach(() => { + spies.forEach((spy) => spy.mockRestore()); + spies = []; + setQuotaAutoBalanceEnabled(false); + updateQuotaCache({}); + (globalThis as { localStorage?: unknown }).localStorage = originalStorage; + }); + + it('computes priority and weight tiers based on bottleneck and average remaining quota', () => { + const healthy: QuotaState = { + status: 'success', + rows: [ + { label: '5-hour window', remainingPercent: 81 }, + { label: '7-day window', remainingPercent: 72 }, + ], + }; + expect(computeQuotaRoutingTarget(healthy)).toEqual({ + priority: AUTO_BALANCE_PRIORITY.healthy, + weight: Math.round(72 * 0.7 + 76.5 * 0.3), + bottleneckPercent: 72, + averagePercent: 76.5, + }); + + const lowWarning: QuotaState = { + status: 'success', + rows: [ + { label: '5-hour window', remainingPercent: 81 }, + { label: '7-day window', remainingPercent: 5 }, + ], + }; + expect(computeQuotaRoutingTarget(lowWarning)).toEqual({ + priority: AUTO_BALANCE_PRIORITY.low, + weight: 5, + bottleneckPercent: 5, + averagePercent: 43, + }); + + const critical: QuotaState = { + status: 'success', + rows: [ + { label: '5-hour window', remainingPercent: 2 }, + { label: '7-day window', remainingPercent: 50 }, + ], + }; + expect(computeQuotaRoutingTarget(critical)).toEqual({ + priority: AUTO_BALANCE_PRIORITY.critical, + weight: 1, + bottleneckPercent: 2, + averagePercent: 26, + }); + }); + + it('ranks depleted credentials below credentials left at the default priority', () => { + expect(AUTO_BALANCE_PRIORITY.healthy).toBe(0); + expect(AUTO_BALANCE_PRIORITY.low).toBeLessThan(0); + expect(AUTO_BALANCE_PRIORITY.critical).toBeLessThan(AUTO_BALANCE_PRIORITY.low); + }); + + it('ignores model-scoped limits when account-wide windows exist', () => { + const target = computeQuotaRoutingTarget({ + status: 'success', + rows: [ + { label: '5-hour window', remainingPercent: 80 }, + { label: '7-day window', remainingPercent: 60 }, + { label: '7-day Opus', remainingPercent: 0, scoped: true }, + ], + }); + expect(target?.priority).toBe(AUTO_BALANCE_PRIORITY.healthy); + expect(target?.bottleneckPercent).toBe(60); + }); + + it('treats the Claude OAuth apps window as account-wide and model windows as scoped', () => { + const rows = quotaRowsFor('claude', { + five_hour: { utilization: 20 }, + seven_day: { utilization: 40 }, + seven_day_oauth_apps: { utilization: 100 }, + seven_day_opus: { utilization: 100 }, + }); + expect(rows.map((row) => Boolean(row.scoped))).toEqual([false, false, false, true]); + const target = computeQuotaRoutingTarget({ status: 'success', rows }); + expect(target?.priority).toBe(AUTO_BALANCE_PRIORITY.critical); + expect(target?.bottleneckPercent).toBe(0); + }); + + it('uses the healthiest model group when only grouped limits exist', () => { + const target = computeQuotaRoutingTarget({ + status: 'success', + rows: [ + { label: 'Claude · 5h', groupLabel: 'Claude', remainingPercent: 0 }, + { label: 'Gemini · 5h', groupLabel: 'Gemini', remainingPercent: 90 }, + { label: 'Gemini · weekly', groupLabel: 'Gemini', remainingPercent: 70 }, + ], + }); + expect(target).toEqual({ + priority: AUTO_BALANCE_PRIORITY.healthy, + weight: Math.round(70 * 0.7 + 35 * 0.3), + bottleneckPercent: 70, + averagePercent: 35, + }); + }); + + it('reads current values from the stored credential and skips redundant PATCH requests', async () => { + const stored = { 'claude-test.json': { priority: 0, weight: 70 } as Record }; + const api = fakeApi(stored); + setQuotaAutoBalanceEnabled(true); + const quota: QuotaState = { + status: 'success', + rows: [ + { label: '5h', remainingPercent: 70 }, + { label: '7d', remainingPercent: 70 }, + ], + }; + + await syncCredentialQuotaRouting({ ...claudeFile }, quota, api); + expect(api.patches).toEqual([]); + }); + + it('patches via commitFetchedQuota without mutating the file and notifies subscribers', async () => { + const stored = { 'claude-test.json': { priority: 3 } as Record }; + const api = fakeApi(stored); + spies.push(spyOn(managementApi, 'get').mockImplementation(api.get as never)); + spies.push(spyOn(managementApi, 'patch').mockImplementation(api.patch as never)); + const changes: Array<[string, unknown]> = []; + const unsubscribe = subscribeQuotaRoutingChanges((name, fields) => changes.push([name, fields])); + setQuotaAutoBalanceEnabled(true); + expect(getQuotaAutoBalanceEnabled()).toBe(true); + + const file = { ...claudeFile }; + commitFetchedQuota(captureQuotaCacheGeneration(), file, { + status: 'success', + rows: [ + { label: '5h', remainingPercent: 70 }, + { label: '7d', remainingPercent: 70 }, + ], + }); + await new Promise((resolve) => setTimeout(resolve, 0)); + unsubscribe(); + + expect(api.patches).toEqual([{ name: 'claude-test.json', priority: 0, weight: 70 }]); + expect(file.priority).toBe(3); + expect(changes).toEqual([['claude-test.json', { priority: 0, weight: 70 }]]); + + const files = [file, { name: 'other.json' }]; + const updated = applyQuotaRoutingChange(files, 'claude-test.json', { priority: 0, weight: 70 }); + expect(updated).not.toBe(files); + expect(updated[0]).toEqual({ ...claudeFile, priority: 0, weight: 70 }); + expect(updated[1]).toBe(files[1]); + }); + + it('restores manual priority and weight captured before the first auto-balance change', async () => { + const stored = { 'claude-test.json': { priority: -10, weight: '' } as Record }; + const api = fakeApi(stored); + const lowQuota: QuotaState = { status: 'success', rows: [{ label: '5h', remainingPercent: 10 }] }; + const healthyQuota: QuotaState = { status: 'success', rows: [{ label: '5h', remainingPercent: 90 }] }; + setQuotaAutoBalanceEnabled(true); + + await syncCredentialQuotaRouting({ ...claudeFile }, lowQuota, api); + await syncCredentialQuotaRouting({ ...claudeFile }, healthyQuota, api); + expect(stored['claude-test.json']).toEqual({ priority: 0, weight: 90 }); + + setQuotaAutoBalanceEnabled(false); + await restoreQuotaRoutingOriginals(api); + expect(stored['claude-test.json']).toEqual({ priority: -10, weight: null }); + + api.patches.length = 0; + await restoreQuotaRoutingOriginals(api); + expect(api.patches).toEqual([]); + }); + + it('does not re-apply auto values when a sync finishes after auto-balance is turned off', async () => { + const stored = { 'claude-test.json': { priority: 7, weight: 3 } as Record }; + const api = fakeApi(stored); + const healthyQuota: QuotaState = { status: 'success', rows: [{ label: '5h', remainingPercent: 90 }] }; + setQuotaAutoBalanceEnabled(true); + await syncCredentialQuotaRouting({ ...claudeFile }, healthyQuota, api); + expect(stored['claude-test.json']).toEqual({ priority: 0, weight: 90 }); + + let releaseDownload = () => {}; + api.beforeDownload = () => new Promise((resolve) => { releaseDownload = resolve; }); + const lateSync = syncCredentialQuotaRouting({ ...claudeFile }, + { status: 'success', rows: [{ label: '5h', remainingPercent: 10 }] }, api); + await new Promise((resolve) => setTimeout(resolve, 0)); + + setQuotaAutoBalanceEnabled(false); + api.beforeDownload = undefined; + const restore = restoreQuotaRoutingOriginals(api); + releaseDownload(); + await Promise.all([lateSync, restore]); + + expect(stored['claude-test.json']).toEqual({ priority: 7, weight: 3 }); + }); + + it('keeps manual edits made while enabled and forgets deleted credentials on restore', async () => { + const stored: Record> = { + 'claude-test.json': { priority: 3 }, + 'gone.json': { priority: 4 }, + }; + const api = fakeApi(stored); + const quota: QuotaState = { status: 'success', rows: [{ label: '5h', remainingPercent: 90 }] }; + setQuotaAutoBalanceEnabled(true); + await syncCredentialQuotaRouting({ ...claudeFile }, quota, api); + await syncCredentialQuotaRouting({ ...claudeFile, name: 'gone.json' }, quota, api); + + stored['claude-test.json'] = { priority: 50, weight: 200 }; + delete stored['gone.json']; + setQuotaAutoBalanceEnabled(false); + api.patches.length = 0; + await restoreQuotaRoutingOriginals(api); + + expect(api.patches).toEqual([]); + expect(stored['claude-test.json']).toEqual({ priority: 50, weight: 200 }); + expect(localStorage.getItem('easy-cli-proxy-api.quota-auto-balance.originals')).toBeNull(); + }); + + it('balances credentials whose unrelated metadata would fail settings validation', async () => { + const stored = { 'claude-test.json': { priority: 3, headers: { 'x-retry': 3 }, excluded_models: 42 } as Record }; + const api = fakeApi(stored); + setQuotaAutoBalanceEnabled(true); + await syncCredentialQuotaRouting({ ...claudeFile }, + { status: 'success', rows: [{ label: '5h', remainingPercent: 90 }] }, api); + expect(api.patches).toEqual([{ name: 'claude-test.json', priority: 0, weight: 90 }]); + }); + + it('follows the toggle when another tab changes it', () => { + const target = new EventTarget(); + const unwatch = watchQuotaAutoBalanceStorage(target); + setQuotaAutoBalanceEnabled(true); + + localStorage.setItem('easy-cli-proxy-api.quota-auto-balance', 'off'); + target.dispatchEvent(Object.assign(new Event('storage'), { key: 'easy-cli-proxy-api.quota-auto-balance' })); + expect(getQuotaAutoBalanceEnabled()).toBe(false); + + localStorage.setItem('easy-cli-proxy-api.quota-auto-balance', 'on'); + target.dispatchEvent(Object.assign(new Event('storage'), { key: 'unrelated' })); + expect(getQuotaAutoBalanceEnabled()).toBe(false); + unwatch(); + }); + + it('skips syncing once another tab has turned auto-balance off', async () => { + const stored = { 'claude-test.json': { priority: 7 } as Record }; + const api = fakeApi(stored); + setQuotaAutoBalanceEnabled(true); + localStorage.setItem('easy-cli-proxy-api.quota-auto-balance', 'off'); + + await syncCredentialQuotaRouting({ ...claudeFile }, + { status: 'success', rows: [{ label: '5h', remainingPercent: 90 }] }, api); + + expect(api.patches).toEqual([]); + expect(getQuotaAutoBalanceEnabled()).toBe(false); + }); +});