@pedrofariasx/qwenproxy 1.3.3 → 1.4.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "@pedrofariasx/qwenproxy",
3
- "version": "1.3.3",
3
+ "version": "1.4.0",
4
4
  "description": "Local OpenAI-compatible proxy API that routes requests to Qwen (chat.qwen.ai) via Playwright browser automation.",
5
5
  "main": "index.js",
6
6
  "scripts": {
@@ -1,4 +1,4 @@
1
- import { QwenAccount, loadAccounts } from './accounts.js'
1
+ import { QwenAccount, loadAccounts, updateAccountCooldown } from './accounts.js'
2
2
  import { config } from './config.js'
3
3
 
4
4
  let currentIndex = 0
@@ -21,6 +21,20 @@ function getCachedAccounts(): QwenAccount[] {
21
21
  if (!accountsCache || (now - accountsCacheTimestamp) > ACCOUNTS_CACHE_TTL) {
22
22
  accountsCache = loadAccounts()
23
23
  accountsCacheTimestamp = now
24
+
25
+ // Sync memory cooldowns map from database values
26
+ for (const account of accountsCache) {
27
+ if (account.cooldown_until && account.cooldown_until > now) {
28
+ cooldowns.set(account.id, {
29
+ until: account.cooldown_until,
30
+ reason: account.cooldown_reason || 'RateLimited',
31
+ })
32
+ } else {
33
+ if (cooldowns.has(account.id)) {
34
+ cooldowns.delete(account.id)
35
+ }
36
+ }
37
+ }
24
38
  }
25
39
  return accountsCache
26
40
  }
@@ -31,15 +45,35 @@ export function invalidateAccountsCache(): void {
31
45
  }
32
46
 
33
47
  export function markAccountRateLimited(accountId: string, cooldownMs?: number, reason?: string): void {
48
+ const duration = cooldownMs ?? DEFAULT_COOLDOWN_MS
49
+ const until = Date.now() + duration
50
+ const cooldownReason = reason ?? 'RateLimited'
51
+
34
52
  cooldowns.set(accountId, {
35
- until: Date.now() + (cooldownMs ?? DEFAULT_COOLDOWN_MS),
36
- reason: reason ?? 'RateLimited',
53
+ until,
54
+ reason: cooldownReason,
37
55
  })
38
- console.log(`[AccountManager] Account ${accountId} marked as rate-limited. Cooldown until ${new Date(Date.now() + (cooldownMs ?? DEFAULT_COOLDOWN_MS)).toISOString()}`)
56
+
57
+ if (accountId !== 'global') {
58
+ try {
59
+ updateAccountCooldown(accountId, until, cooldownReason)
60
+ } catch (err) {
61
+ console.error(`[AccountManager] Failed to save cooldown to DB for account ${accountId}:`, (err as Error).message)
62
+ }
63
+ }
64
+
65
+ console.log(`[AccountManager] Account ${accountId} marked as rate-limited. Cooldown until ${new Date(until).toISOString()}`)
39
66
  }
40
67
 
41
68
  export function clearAccountCooldown(accountId: string): void {
42
69
  cooldowns.delete(accountId)
70
+ if (accountId !== 'global') {
71
+ try {
72
+ updateAccountCooldown(accountId, 0, null)
73
+ } catch (err) {
74
+ console.error(`[AccountManager] Failed to clear cooldown in DB for account ${accountId}:`, (err as Error).message)
75
+ }
76
+ }
43
77
  }
44
78
 
45
79
  export function getAccountCooldownInfo(accountId: string): { onCooldown: boolean; remainingMs: number; reason: string } | null {
@@ -48,6 +82,13 @@ export function getAccountCooldownInfo(accountId: string): { onCooldown: boolean
48
82
  const remaining = entry.until - Date.now()
49
83
  if (remaining <= 0) {
50
84
  cooldowns.delete(accountId)
85
+ if (accountId !== 'global') {
86
+ try {
87
+ updateAccountCooldown(accountId, 0, null)
88
+ } catch (err) {
89
+ console.error(`[AccountManager] Failed to clear expired cooldown in DB:`, (err as Error).message)
90
+ }
91
+ }
51
92
  return null
52
93
  }
53
94
  return { onCooldown: true, remainingMs: remaining, reason: entry.reason }
@@ -88,25 +129,33 @@ export function getNextAccount(forceReset?: boolean): QwenAccount | null {
88
129
  return best
89
130
  }
90
131
 
91
- export function getNextAvailableAccount(skipAccountId?: string): QwenAccount | null {
132
+ export function getNextAvailableAccount(triedAccountIds?: Set<string> | string): QwenAccount | null {
92
133
  const accounts = getCachedAccounts()
93
134
  if (accounts.length === 0) return null
94
135
 
136
+ let triedSet: Set<string>
137
+ if (triedAccountIds instanceof Set) {
138
+ triedSet = triedAccountIds
139
+ } else {
140
+ triedSet = new Set(triedAccountIds ? [triedAccountIds] : [])
141
+ }
142
+
143
+ // 1. Try to find an untried account that is NOT on cooldown
95
144
  for (let i = 0; i < accounts.length; i++) {
96
145
  const idx = (currentIndex + i) % accounts.length
97
146
  const account = accounts[idx]
98
- if (skipAccountId && account.id === skipAccountId) continue
147
+ if (triedSet.has(account.id)) continue
99
148
  if (!isAccountOnCooldown(account.id)) {
100
149
  currentIndex = (idx + 1) % accounts.length
101
150
  return account
102
151
  }
103
152
  }
104
153
 
105
- // All remaining accounts on cooldown return the one with shortest cooldown
154
+ // 2. If all untried accounts are on cooldown, return the untried one with the shortest remaining cooldown
106
155
  let best: QwenAccount | null = null
107
156
  let bestRemaining = Infinity
108
157
  for (const account of accounts) {
109
- if (skipAccountId && account.id === skipAccountId) continue
158
+ if (triedSet.has(account.id)) continue
110
159
  const info = getAccountCooldownInfo(account.id)
111
160
  if (info && info.remainingMs < bestRemaining) {
112
161
  bestRemaining = info.remainingMs
@@ -6,6 +6,8 @@ export interface QwenAccount {
6
6
  id: string
7
7
  email: string
8
8
  password: string
9
+ cooldown_until?: number
10
+ cooldown_reason?: string | null
9
11
  }
10
12
 
11
13
  let accountsCache: QwenAccount[] | null = null
@@ -16,7 +18,7 @@ function getCachedAccounts(): QwenAccount[] {
16
18
  const now = Date.now()
17
19
  if (!accountsCache || (now - accountsCacheTimestamp) > ACCOUNTS_CACHE_TTL) {
18
20
  const db = getDatabase()
19
- accountsCache = db.prepare('SELECT id, email, password FROM accounts ORDER BY created_at ASC').all() as QwenAccount[]
21
+ accountsCache = db.prepare('SELECT id, email, password, cooldown_until, cooldown_reason FROM accounts ORDER BY created_at ASC').all() as QwenAccount[]
20
22
  accountsCacheTimestamp = now
21
23
  }
22
24
  return accountsCache
@@ -73,6 +75,12 @@ export function listAccounts(): QwenAccount[] {
73
75
 
74
76
  export function getAccountCredentials(id: string): QwenAccount | undefined {
75
77
  const db = getDatabase()
76
- const row = db.prepare('SELECT id, email, password FROM accounts WHERE id = ?').get(id)
78
+ const row = db.prepare('SELECT id, email, password, cooldown_until, cooldown_reason FROM accounts WHERE id = ?').get(id)
77
79
  return row as QwenAccount | undefined
78
80
  }
81
+
82
+ export function updateAccountCooldown(id: string, cooldownUntil: number, reason: string | null): void {
83
+ const db = getDatabase()
84
+ db.prepare('UPDATE accounts SET cooldown_until = ?, cooldown_reason = ? WHERE id = ?').run(cooldownUntil, reason, id)
85
+ invalidateAccountsCache()
86
+ }
@@ -42,6 +42,18 @@ function runMigrations(db: Database.Database): void {
42
42
 
43
43
  CREATE INDEX IF NOT EXISTS idx_accounts_email ON accounts(email);
44
44
  `)
45
+
46
+ // Add cooldown columns if they don't exist
47
+ try {
48
+ db.exec(`ALTER TABLE accounts ADD COLUMN cooldown_until INTEGER DEFAULT 0;`)
49
+ } catch (err) {
50
+ // Column already exists or error
51
+ }
52
+ try {
53
+ db.exec(`ALTER TABLE accounts ADD COLUMN cooldown_reason TEXT;`)
54
+ } catch (err) {
55
+ // Column already exists or error
56
+ }
45
57
  }
46
58
 
47
59
  /**
@@ -25,45 +25,47 @@ const modelContextWindows: Record<string, number> = {
25
25
  }
26
26
 
27
27
  const modelTokenDivisors: Record<string, number> = {
28
- 'qwen3.7-max': 2.2,
29
- 'qwen3.6-max-preview': 2.2,
30
- 'qwen3.5-max-2026-03-08': 2.2,
31
- 'qwen3-max-2026-01-23': 2.2,
32
- 'qwen-latest-series-invite-beta-v24': 2.2,
33
- 'qwen3.7-plus': 2.0,
34
- 'qwen3.6-plus': 2.0,
35
- 'qwen3.6-plus-preview': 2.0,
36
- 'qwen3.5-plus': 2.0,
37
- 'qwen-plus-2025-07-28': 2.0,
38
- 'qwen-latest-series-invite-beta-v16': 2.0,
39
- 'qwen3.5-flash': 1.8,
40
- 'qwen3.5-omni-plus': 1.8,
41
- 'qwen3.5-omni-flash': 1.7,
42
- 'qwen3-omni-flash-2025-12-01': 1.7,
43
- 'qwen3.5-397b-a17b': 1.9,
44
- 'qwen3.5-122b-a10b': 1.9,
45
- 'qwen3.6-35b-a3b': 1.9,
46
- 'qwen3.5-35b-a3b': 1.9,
47
- 'qwen3.6-27b': 1.9,
48
- 'qwen3.5-27b': 1.9,
49
- 'qwen3-coder-plus': 2.3,
50
- 'qwen3-vl-plus': 2.1,
28
+ 'qwen3.7-max': 3.5,
29
+ 'qwen3.6-max-preview': 3.5,
30
+ 'qwen3.5-max-2026-03-08': 3.5,
31
+ 'qwen3-max-2026-01-23': 3.5,
32
+ 'qwen-latest-series-invite-beta-v24': 3.5,
33
+ 'qwen3.7-plus': 3.5,
34
+ 'qwen3.6-plus': 3.5,
35
+ 'qwen3.6-plus-preview': 3.5,
36
+ 'qwen3.5-plus': 3.5,
37
+ 'qwen-plus-2025-07-28': 3.5,
38
+ 'qwen-latest-series-invite-beta-v16': 3.5,
39
+ 'qwen3.5-flash': 3.2,
40
+ 'qwen3.5-omni-plus': 3.0,
41
+ 'qwen3.5-omni-flash': 3.0,
42
+ 'qwen3-omni-flash-2025-12-01': 3.0,
43
+ 'qwen3.5-397b-a17b': 3.2,
44
+ 'qwen3.5-122b-a10b': 3.2,
45
+ 'qwen3.6-35b-a3b': 3.2,
46
+ 'qwen3.5-35b-a3b': 3.2,
47
+ 'qwen3.6-27b': 3.2,
48
+ 'qwen3.5-27b': 3.2,
49
+ 'qwen3-coder-plus': 3.8,
50
+ 'qwen3-vl-plus': 3.5,
51
51
  }
52
52
 
53
53
  const defaultContextWindow = 131072
54
- const defaultTokenDivisor = 2.0
55
- export const MAX_PAYLOAD_SIZE = 10 * 1024 * 1024
54
+ const defaultTokenDivisor = 3.5
55
+ export const MAX_PAYLOAD_SIZE = 50 * 1024 * 1024
56
56
 
57
57
  export function setModelContextWindow(modelId: string, contextWindow: number): void {
58
58
  modelContextWindows[modelId] = contextWindow
59
59
  }
60
60
 
61
- export function getModelContextWindow(modelId: string): number {
61
+ export function getModelContextWindow(modelId?: string): number {
62
+ if (!modelId) return defaultContextWindow
62
63
  const baseId = modelId.replace('-no-thinking', '')
63
64
  return modelContextWindows[baseId] ?? defaultContextWindow
64
65
  }
65
66
 
66
- export function getModelTokenDivisor(modelId: string): number {
67
+ export function getModelTokenDivisor(modelId?: string): number {
68
+ if (!modelId) return defaultTokenDivisor
67
69
  const baseId = modelId.replace('-no-thinking', '')
68
70
  return modelTokenDivisors[baseId] ?? defaultTokenDivisor
69
71
  }
@@ -21,6 +21,7 @@ import { QwenStreamParser, ParsedChunkResult } from '../utils/qwen-stream-parser
21
21
  import { getModelContextWindow } from '../core/model-registry.js'
22
22
  import { truncateMessages, estimateTokenCount } from '../utils/context-truncation.js';
23
23
  import { getNextAccount, getNextAvailableAccount, markAccountRateLimited, getAccountCooldownInfo } from '../core/account-manager.js';
24
+ import { loadAccounts } from '../core/accounts.js';
24
25
  import { registerStream, removeStream, getStream } from '../core/stream-registry.js';
25
26
  import { metrics } from '../core/metrics.js'
26
27
 
@@ -284,7 +285,7 @@ export async function chatCompletions(c: Context) {
284
285
  const accountEmail = account.email;
285
286
 
286
287
  if (triedAccountIds.has(accountId)) {
287
- account = getNextAvailableAccount(accountId);
288
+ account = getNextAvailableAccount(triedAccountIds);
288
289
  continue;
289
290
  }
290
291
  triedAccountIds.add(accountId);
@@ -292,7 +293,7 @@ export async function chatCompletions(c: Context) {
292
293
  const cooldownInfo = getAccountCooldownInfo(accountId);
293
294
  if (cooldownInfo && accountId !== 'global') {
294
295
  console.log(`[Chat] Skipping account ${accountEmail} (${accountId}) — on cooldown for ${Math.round(cooldownInfo.remainingMs / 1000)}s (${cooldownInfo.reason})`);
295
- account = getNextAvailableAccount(accountId);
296
+ account = getNextAvailableAccount(triedAccountIds);
296
297
  continue;
297
298
  }
298
299
 
@@ -329,9 +330,10 @@ export async function chatCompletions(c: Context) {
329
330
 
330
331
  if (err.upstreamCode === 'RateLimited' || err.upstreamStatus === 429) {
331
332
  const hourHint = err.message?.match(/Wait about (\d+) hour/);
332
- const cooldownMs = hourHint ? parseInt(hourHint[1]) * 60 * 60 * 1000 : undefined;
333
+ const hours = hourHint ? parseInt(hourHint[1]) : 24;
334
+ const cooldownMs = hours * 60 * 60 * 1000;
333
335
  markAccountRateLimited(accountId, cooldownMs, 'RateLimited');
334
- console.warn(`[Chat] Account ${accountEmail} (${accountId}) rate-limited. Marked for cooldown.`);
336
+ console.warn(`[Chat] Account ${accountEmail} (${accountId}) rate-limited. Entering cooldown for ${hours} hours.`);
335
337
  lastError = err;
336
338
  break;
337
339
  }
@@ -364,11 +366,17 @@ export async function chatCompletions(c: Context) {
364
366
  break;
365
367
  }
366
368
 
367
- account = getNextAvailableAccount(accountId);
369
+ account = getNextAvailableAccount(triedAccountIds);
368
370
  }
369
371
 
370
372
  if (!stream) {
371
373
  removeStream(completionId);
374
+ // Check if all accounts are on cooldown
375
+ const accounts = loadAccounts();
376
+ const allOnCooldown = accounts.every(a => getAccountCooldownInfo(a.id) !== null);
377
+ if (allOnCooldown) {
378
+ console.warn(`[Chat] CRITICAL: All ${accounts.length} accounts are currently rate-limited or on cooldown!`);
379
+ }
372
380
  throw lastError || new Error('All accounts failed');
373
381
  }
374
382
 
@@ -188,9 +188,28 @@ export async function getBasicHeaders(accountId?: string): Promise<{ cookie: str
188
188
  }
189
189
 
190
190
  const cache = getAccountHeaderCache(cacheKey);
191
+ let bxUa = cache.currentHeaders['bx-ua'];
192
+ let bxUmidtoken = cache.currentHeaders['bx-umidtoken'];
191
193
  const bxV = cache.currentHeaders['bx-v'] || '2.5.36';
192
- const bxUa = cache.currentHeaders['bx-ua'];
193
- const bxUmidtoken = cache.currentHeaders['bx-umidtoken'];
194
+
195
+ // Auto-recover missing anti-fraud headers by triggering full header interception
196
+ if (!bxUa || !bxUmidtoken) {
197
+ console.log(`[Playwright] Missing bx-ua/bx-umidtoken for ${cacheKey}, triggering header interception...`);
198
+ try {
199
+ const result = await getQwenHeaders(true, accountId);
200
+ bxUa = result.headers['bx-ua'];
201
+ bxUmidtoken = result.headers['bx-umidtoken'];
202
+ return {
203
+ cookie: await getCookies(accountId),
204
+ userAgent,
205
+ bxV: result.headers['bx-v'] || bxV,
206
+ bxUa,
207
+ bxUmidtoken,
208
+ };
209
+ } catch (err: any) {
210
+ console.warn(`[Playwright] Failed to auto-recover headers for ${cacheKey}: ${err.message}`);
211
+ }
212
+ }
194
213
 
195
214
  return { cookie, userAgent, bxV, bxUa, bxUmidtoken };
196
215
  }
@@ -649,6 +668,25 @@ export async function initPlaywrightForAccount(account: QwenAccount, headless =
649
668
  if (!hasAuthCookie && account.email && account.password) {
650
669
  await loginToQwenWithContext(acctContext, acctPage, account.email, account.password);
651
670
  }
671
+
672
+ // Navigate to Qwen home to validate session and populate cookies
673
+ try {
674
+ await acctPage.goto('https://chat.qwen.ai/', { waitUntil: 'domcontentloaded', timeout: 15000 });
675
+ const url = acctPage.url();
676
+ if (url.includes('auth') || url.includes('login')) {
677
+ if (account.email && account.password) {
678
+ console.log(`[Playwright] Session expired for ${account.email}, re-logging in...`);
679
+ await loginToQwenWithContext(acctContext, acctPage, account.email, account.password);
680
+ await acctPage.goto('https://chat.qwen.ai/', { waitUntil: 'domcontentloaded', timeout: 15000 });
681
+ } else {
682
+ console.warn(`[Playwright] Session expired for account ${account.id} but no credentials available for re-login.`);
683
+ }
684
+ } else {
685
+ console.log(`[Playwright] Session validated for ${account.email}.`);
686
+ }
687
+ } catch (err: any) {
688
+ console.warn(`[Playwright] Failed to validate session for ${account.email}: ${err.message}`);
689
+ }
652
690
  }
653
691
 
654
692
  export async function launchManualLoginAccount(accountId: string, browserType: BrowserType = 'chromium'): Promise<{ context: BrowserContext, page: Page }> {
@@ -1,5 +1,6 @@
1
1
  import { getQwenHeaders, getBasicHeaders } from './playwright.js';
2
2
  import { MAX_PAYLOAD_SIZE } from '../core/model-registry.js';
3
+ import { markAccountRateLimited } from '../core/account-manager.js';
3
4
  import crypto from 'crypto';
4
5
 
5
6
  const CACHED_TIMEZONE = new Date().toString().split(' (')[0];
@@ -133,8 +134,32 @@ async function createRealQwenChat(header: Record<string, string>): Promise<strin
133
134
  signal: AbortSignal.timeout(30000),
134
135
  });
135
136
 
136
- if (!response.ok) throw new Error(`Failed to create chat: ${response.status}`);
137
+ if (!response.ok) {
138
+ const errText = await response.text().catch(() => '');
139
+ if (response.status === 429) {
140
+ throw new QwenUpstreamError(
141
+ 'Qwen upstream error: RateLimited: Too many requests.',
142
+ 'RateLimited',
143
+ 429
144
+ );
145
+ }
146
+ throw new Error(`Failed to create chat: ${response.status} - ${errText}`);
147
+ }
137
148
  const json = await response.json();
149
+ if (json && json.success === false) {
150
+ const code = json.data?.code || json.code || 'UpstreamError';
151
+ const details = json.data?.details || json.message || 'Qwen returned an error';
152
+ const wait = json.data?.num !== undefined
153
+ ? ` Wait about ${json.data.num} hour(s) before trying again.`
154
+ : '';
155
+ let status = 502;
156
+ if (code === 'RateLimited') status = 429;
157
+ throw new QwenUpstreamError(
158
+ `Qwen upstream error: ${code}: ${details}.${wait}`,
159
+ code,
160
+ status
161
+ );
162
+ }
138
163
  const chatId = json.chat_id || json.id || json.data?.chat_id || json.data?.id;
139
164
  if (!chatId) throw new Error(`Unexpected chat response: ${JSON.stringify(json).slice(0, 200)}`);
140
165
  return chatId;
@@ -160,7 +185,15 @@ async function refillPoolForAccount(accountId: string) {
160
185
  try {
161
186
  const chatId = await createRealQwenChat(headers);
162
187
  return { chatId, headers, accountId, timestamp: Date.now() };
163
- } catch (err) {
188
+ } catch (err: any) {
189
+ if (err instanceof QwenUpstreamError) {
190
+ if (err.upstreamCode === 'RateLimited' || err.upstreamStatus === 429) {
191
+ const hourHint = err.message?.match(/Wait about (\d+) hour/);
192
+ const cooldownMs = hourHint ? parseInt(hourHint[1]) * 60 * 60 * 1000 : undefined;
193
+ markAccountRateLimited(accountId, cooldownMs, 'RateLimited');
194
+ console.warn(`[WarmPool] Account ${accountId} rate-limited during chat creation. Marked for cooldown.`);
195
+ }
196
+ }
164
197
  console.error(`[WarmPool] chat creation failed for ${accountId}:`, (err as Error).message);
165
198
  return null;
166
199
  }
@@ -564,6 +597,26 @@ export async function createQwenStream(
564
597
  }
565
598
 
566
599
  if (retryResponse.ok && retryResponse.body) {
600
+ try {
601
+ const errorJson = JSON.parse(retryPeek);
602
+ if (errorJson && (errorJson.success === false || errorJson.error)) {
603
+ const code = errorJson.data?.code || errorJson.code || 'UpstreamError';
604
+ const details = errorJson.data?.details || errorJson.message || errorJson.error?.message || 'Qwen returned an error';
605
+ const wait = errorJson.data?.num !== undefined
606
+ ? ` Wait about ${errorJson.data.num} hour(s) before trying again.`
607
+ : '';
608
+ let status = 502;
609
+ if (code === 'RateLimited') status = 429;
610
+
611
+ throw new QwenUpstreamError(
612
+ `Qwen upstream error: ${code}: ${details}.${wait}`,
613
+ code,
614
+ status,
615
+ );
616
+ }
617
+ } catch (e) {
618
+ if (e instanceof QwenUpstreamError) throw e;
619
+ }
567
620
  return { stream: retryResponse.body, headers: freshHeaders, uiSessionId: chatId, controller: retryController, accountId: chatEntry.accountId };
568
621
  }
569
622
  } catch (retryErr) {
@@ -576,6 +629,27 @@ export async function createQwenStream(
576
629
  'FAIL_SYS_USER_VALIDATE',
577
630
  403,
578
631
  );
632
+ } else {
633
+ try {
634
+ const errorJson = JSON.parse(peekText);
635
+ if (errorJson && (errorJson.success === false || errorJson.error)) {
636
+ const code = errorJson.data?.code || errorJson.code || 'UpstreamError';
637
+ const details = errorJson.data?.details || errorJson.message || errorJson.error?.message || 'Qwen returned an error';
638
+ const wait = errorJson.data?.num !== undefined
639
+ ? ` Wait about ${errorJson.data.num} hour(s) before trying again.`
640
+ : '';
641
+ let status = 502;
642
+ if (code === 'RateLimited') status = 429;
643
+
644
+ throw new QwenUpstreamError(
645
+ `Qwen upstream error: ${code}: ${details}.${wait}`,
646
+ code,
647
+ status,
648
+ );
649
+ }
650
+ } catch (e) {
651
+ if (e instanceof QwenUpstreamError) throw e;
652
+ }
579
653
  }
580
654
  }
581
655
 
@@ -6,20 +6,20 @@ test('estimateTokenCount: returns 0 for empty string', () => {
6
6
  assert.strictEqual(estimateTokenCount(''), 0);
7
7
  });
8
8
 
9
- test('estimateTokenCount: estimates tokens conservatively using 2.5 divisor', () => {
9
+ test('estimateTokenCount: estimates tokens conservatively using default divisor', () => {
10
10
  assert.strictEqual(estimateTokenCount('hello'), 2);
11
- assert.strictEqual(estimateTokenCount('a'.repeat(100)), 40);
12
- assert.strictEqual(estimateTokenCount('a'.repeat(250)), 100);
13
- assert.strictEqual(estimateTokenCount('a'.repeat(2500)), 1000);
11
+ assert.strictEqual(estimateTokenCount('a'.repeat(100)), 29);
12
+ assert.strictEqual(estimateTokenCount('a'.repeat(250)), 72);
13
+ assert.strictEqual(estimateTokenCount('a'.repeat(2500)), 715);
14
14
  });
15
15
 
16
16
  test('estimateTokenCount: handles single character', () => {
17
17
  assert.strictEqual(estimateTokenCount('x'), 1);
18
18
  });
19
19
 
20
- test('estimateTokenCount: rounds up for non-multiples of 2.5', () => {
20
+ test('estimateTokenCount: rounds up for non-multiples of default divisor', () => {
21
21
  assert.strictEqual(estimateTokenCount('ab'), 1);
22
- assert.strictEqual(estimateTokenCount('abc'), 2);
22
+ assert.strictEqual(estimateTokenCount('abc'), 1);
23
23
  assert.strictEqual(estimateTokenCount('abcd'), 2);
24
24
  });
25
25
 
@@ -1,6 +1,6 @@
1
1
  import { test } from 'node:test';
2
2
  import assert from 'node:assert';
3
- import { getNextAccount, invalidateAccountsCache } from '../core/account-manager.ts';
3
+ import { getNextAccount, getNextAvailableAccount, markAccountRateLimited, clearAccountCooldown, invalidateAccountsCache } from '../core/account-manager.ts';
4
4
  import { addAccount, removeAccount, loadAccounts } from '../core/accounts.ts';
5
5
 
6
6
  test('Account Rotation: Round-Robin rotation cycle', async () => {
@@ -29,10 +29,15 @@ test('Account Rotation: Round-Robin rotation cycle', async () => {
29
29
  assert.ok(third);
30
30
  assert.ok(fourth);
31
31
 
32
- assert.strictEqual(first.email, 'account1@test.com');
33
- assert.strictEqual(second.email, 'account2@test.com');
34
- assert.strictEqual(third.email, 'account3@test.com');
35
- assert.strictEqual(fourth.email, 'account1@test.com');
32
+ const allAccounts = loadAccounts();
33
+ const firstIdx = allAccounts.findIndex(a => a.id === first.id);
34
+ const secondIdx = allAccounts.findIndex(a => a.id === second.id);
35
+ const thirdIdx = allAccounts.findIndex(a => a.id === third.id);
36
+ const fourthIdx = allAccounts.findIndex(a => a.id === fourth.id);
37
+
38
+ assert.strictEqual(secondIdx, (firstIdx + 1) % allAccounts.length);
39
+ assert.strictEqual(thirdIdx, (secondIdx + 1) % allAccounts.length);
40
+ assert.strictEqual(fourthIdx, (thirdIdx + 1) % allAccounts.length);
36
41
  } finally {
37
42
  const current = loadAccounts();
38
43
  for (const acc of current) {
@@ -43,3 +48,57 @@ test('Account Rotation: Round-Robin rotation cycle', async () => {
43
48
  invalidateAccountsCache();
44
49
  }
45
50
  });
51
+
52
+ test('Account Cooldown: Database persistence and recovery', async () => {
53
+ const email = 'cooldown-test@test.com';
54
+ let accountId = '';
55
+
56
+ try {
57
+ const newAcct = addAccount(email, 'password123');
58
+ accountId = newAcct.id;
59
+ invalidateAccountsCache();
60
+
61
+ // Mark as rate-limited with a 1-hour cooldown
62
+ const cooldownMs = 60 * 60 * 1000;
63
+ markAccountRateLimited(accountId, cooldownMs, 'RateLimited');
64
+
65
+ // Force reloading accounts from DB (simulating restart)
66
+ invalidateAccountsCache();
67
+
68
+ // Check if the loaded account has the cooldown synced from DB
69
+ const loadedAccounts = loadAccounts();
70
+ const target = loadedAccounts.find(a => a.id === accountId);
71
+ assert.ok(target);
72
+ assert.ok(target.cooldown_until);
73
+ assert.ok(target.cooldown_until > Date.now());
74
+ assert.strictEqual(target.cooldown_reason, 'RateLimited');
75
+
76
+ // Verify rotation skips it
77
+ const triedSet = new Set<string>();
78
+ triedSet.add('dummy-id'); // to force getNextAvailableAccount check
79
+ const available = getNextAvailableAccount(triedSet);
80
+ // Since our test account is on cooldown, if it was returned, it means no other account was available,
81
+ // or if we have other non-cooldown accounts, it returned one of them.
82
+ if (available && available.id === accountId) {
83
+ // If it returned our test account, it must be because all accounts are on cooldown.
84
+ // Let's assert that the cooldown is actually registered in memory.
85
+ const info = getNextAccount();
86
+ // It shouldn't be the first option if others are available
87
+ }
88
+
89
+ // Clear cooldown and verify it is updated in DB
90
+ clearAccountCooldown(accountId);
91
+ invalidateAccountsCache();
92
+
93
+ const reloaded = loadAccounts().find(a => a.id === accountId);
94
+ assert.ok(reloaded);
95
+ assert.strictEqual(reloaded.cooldown_until || 0, 0);
96
+ assert.strictEqual(reloaded.cooldown_reason, null);
97
+
98
+ } finally {
99
+ if (accountId) {
100
+ removeAccount(accountId);
101
+ }
102
+ invalidateAccountsCache();
103
+ }
104
+ });
@@ -1,7 +1,7 @@
1
1
  import { getModelTokenDivisor } from '../core/model-registry.js'
2
2
 
3
3
  export function estimateTokenCount(text: string, modelId?: string): number {
4
- const divisor = modelId ? getModelTokenDivisor(modelId) : 2.0
4
+ const divisor = getModelTokenDivisor(modelId)
5
5
  return Math.ceil(text.length / divisor)
6
6
  }
7
7
 
@@ -36,7 +36,7 @@ export function truncateMessages(
36
36
  systemPrompt: string = '',
37
37
  modelId?: string
38
38
  ): Array<{ role: string; content: string }> {
39
- const divisor = modelId ? getModelTokenDivisor(modelId) : 2.0
39
+ const divisor = getModelTokenDivisor(modelId)
40
40
  const systemTokens = estimateTokenCount(systemPrompt, modelId);
41
41
  const availableTokens = maxContextLength - systemTokens - 500;
42
42