From 40f32ebbf06937b85f7e151169ae93399a1dca04 Mon Sep 17 00:00:00 2001 From: Tam Nhu Tran Date: Sun, 7 Jun 2026 15:32:40 -0400 Subject: [PATCH] feat(quota): add per-provider tier-lock account selection Make tier_lock a per-provider map and honor it in findHealthyAccount and preflightCheck for the selected provider only, so locking one provider never disables failover for others. Add POST /api/accounts/tier-lock with tier validation against known tiers, and serialize quota_management on config write so the lock persists. --- src/cliproxy/quota/quota-manager.ts | 35 +- src/config/loader/yaml-serializer.ts | 19 + src/config/schemas/quota.ts | 37 +- src/web-server/routes/account-routes.ts | 118 ++++- .../cliproxy/quota-manager-tier-lock.test.ts | 431 ++++++++++++++++++ .../unit/web-server/tier-lock-routes.test.ts | 200 ++++++++ 6 files changed, 836 insertions(+), 4 deletions(-) create mode 100644 tests/unit/cliproxy/quota-manager-tier-lock.test.ts create mode 100644 tests/unit/web-server/tier-lock-routes.test.ts diff --git a/src/cliproxy/quota/quota-manager.ts b/src/cliproxy/quota/quota-manager.ts index 8574b079..fcef7c44 100644 --- a/src/cliproxy/quota/quota-manager.ts +++ b/src/cliproxy/quota/quota-manager.ts @@ -35,6 +35,7 @@ import { import type { RuntimeMonitorConfig } from '../../config/unified-config-types'; import { loadOrCreateUnifiedConfig } from '../../config/config-loader-facade'; +import { getTierLockForProvider } from '../../config/schemas/quota'; export type ManagedQuotaProvider = 'agy' | 'claude' | 'codex' | 'gemini' | 'ghcp'; type ManagedQuotaResult = @@ -428,10 +429,20 @@ export async function findHealthyAccount( const accounts = getProviderAccounts(provider); + // When a tier is locked for this specific provider, restrict candidates to + // that tier only. Locks are per-provider so locking "agy" to "ultra" does + // NOT affect failover for "claude", "codex", "gemini", or "ghcp". + // This is intentionally strict: no cross-tier fallback while a lock is active, + // so the user's explicit tier choice is always honored. + const tierLock = getTierLockForProvider(config.quota_management?.manual, provider); + // Filter available accounts const available = accounts.filter( (a) => - !exclude.includes(a.id) && !isAccountPaused(provider, a.id) && !isOnCooldown(provider, a.id) + !exclude.includes(a.id) && + !isAccountPaused(provider, a.id) && + !isOnCooldown(provider, a.id) && + (tierLock === null || (a.tier || 'unknown') === tierLock) ); if (available.length === 0) return null; @@ -623,6 +634,28 @@ export async function preflightCheck(provider: CLIProxyProvider): Promise | null; } /** @@ -98,6 +110,27 @@ export const DEFAULT_MANUAL_QUOTA_CONFIG: ManualQuotaConfig = { tier_lock: null, }; +/** + * Read the tier lock for a specific provider from a ManualQuotaConfig. + * + * Handles three cases: + * 1. tier_lock is null/undefined → no lock active for any provider + * 2. tier_lock is a Record (new shape) → return the value for this provider + * 3. tier_lock is a bare string (legacy shape, pre-per-provider) → treated as + * no lock to avoid silently applying an old global lock to every provider + * + * @returns The tier string to lock to, or null (no lock). + */ +export function getTierLockForProvider( + manual: ManualQuotaConfig | undefined | null, + provider: string +): string | null { + const tierLock = manual?.tier_lock; + if (!tierLock || typeof tierLock !== 'object') return null; + const value = (tierLock as Record)[provider]; + return value ?? null; +} + /** * Default runtime monitor configuration. */ diff --git a/src/web-server/routes/account-routes.ts b/src/web-server/routes/account-routes.ts index 50294ae8..4a6520db 100644 --- a/src/web-server/routes/account-routes.ts +++ b/src/web-server/routes/account-routes.ts @@ -38,7 +38,20 @@ import { } from './account-route-helpers'; import type { AccountConfig } from '../../config/unified-config-types'; import { resolveConfiguredPlainCcsResumeLane } from '../../auth/resume-lane-diagnostics'; -import { isUnifiedMode, loadOrCreateUnifiedConfig } from '../../config/config-loader-facade'; +import { + isUnifiedMode, + loadOrCreateUnifiedConfig, + mutateConfig, +} from '../../config/config-loader-facade'; +import type { AccountTier } from '../../cliproxy/accounts/types'; + +/** Valid account tier values for tier-lock validation */ +const VALID_ACCOUNT_TIERS: ReadonlySet = new Set([ + 'free', + 'pro', + 'ultra', + 'unknown', +]); const router = Router(); @@ -633,4 +646,107 @@ router.post('/solo', async (req: Request, res: Response): Promise => { } }); +/** + * POST /api/accounts/tier-lock - Set or clear the tier_lock for a provider + * + * Body: { provider: string, tier: string | null } + * - tier: the tier name to lock to (e.g. "ultra", "pro"), or null to clear. + * + * Persists via the existing config write path (mutateConfig → quota_management.manual.tier_lock). + * The quota-manager reads this on every preflight/findHealthyAccount call. + */ +router.post('/tier-lock', (req: Request, res: Response): void => { + try { + const { provider, tier } = req.body as { provider?: unknown; tier?: unknown }; + + if (!provider || typeof provider !== 'string') { + res.status(400).json({ error: 'Missing required field: provider (string)' }); + return; + } + + if (!isCLIProxyProvider(provider)) { + res.status(400).json({ error: `Invalid provider: ${provider}` }); + return; + } + + // tier must be explicitly present in body (key exists), and must be string or null + if (!('tier' in req.body)) { + res.status(400).json({ error: 'Missing required field: tier (string | null)' }); + return; + } + + if (tier !== null && typeof tier !== 'string') { + res.status(400).json({ error: 'Invalid tier: must be a string or null' }); + return; + } + + // Validate tier against known AccountTier values (null means clear the lock) + if (tier !== null && !VALID_ACCOUNT_TIERS.has(tier)) { + res.status(400).json({ + error: `Invalid tier: "${tier}". Must be one of: ${[...VALID_ACCOUNT_TIERS].join(', ')}`, + }); + return; + } + + // Persist via existing config write path. + // tier_lock is a per-provider map: { [provider]: tier | null } + // Setting a provider's entry to null clears the lock for that provider only. + mutateConfig((config) => { + if (!config.quota_management) { + config.quota_management = { + mode: 'hybrid', + auto: { + preflight_check: true, + exhaustion_threshold: 5, + tier_priority: ['ultra', 'pro', 'free'], + cooldown_minutes: 5, + }, + manual: { + paused_accounts: [], + forced_default: null, + tier_lock: { [provider]: tier ?? null }, + }, + runtime_monitor: { + enabled: true, + normal_interval_seconds: 300, + critical_interval_seconds: 60, + warn_threshold: 20, + exhaustion_threshold: 5, + cooldown_minutes: 5, + }, + }; + } else { + if (!config.quota_management.manual) { + config.quota_management.manual = { + paused_accounts: [], + forced_default: null, + tier_lock: { [provider]: tier ?? null }, + }; + } else { + // Ensure config.quota_management.manual is an owned object, not a shared + // reference from DEFAULT_MANUAL_QUOTA_CONFIG (which createEmptyUnifiedConfig + // sets via shallow spread). Replacing the reference prevents mutation of + // the module-level default constant. + config.quota_management.manual = { ...config.quota_management.manual }; + + // Ensure tier_lock is a map (guard against legacy string shape or null) + const existing = config.quota_management.manual.tier_lock; + if (!existing || typeof existing !== 'object') { + config.quota_management.manual.tier_lock = { [provider]: tier ?? null }; + } else { + // Spread to own the map too before mutating it + const ownedMap = { ...(existing as Record) }; + ownedMap[provider] = tier ?? null; + config.quota_management.manual.tier_lock = ownedMap; + } + } + } + }); + + res.json({ provider, tier_lock: tier ?? null }); + } catch (error) { + res.status(500).json({ error: (error as Error).message }); + } +}); + export default router; diff --git a/tests/unit/cliproxy/quota-manager-tier-lock.test.ts b/tests/unit/cliproxy/quota-manager-tier-lock.test.ts new file mode 100644 index 00000000..0ba6a66a --- /dev/null +++ b/tests/unit/cliproxy/quota-manager-tier-lock.test.ts @@ -0,0 +1,431 @@ +/** + * Phase 2: quota-manager tier_lock selection tests + * + * Load-bearing requirement: when manual.tier_lock is set in config, + * findHealthyAccount must only return accounts matching that tier. + * Clearing tier_lock (null) restores normal tier-priority selection. + * Existing failover behavior must be unaffected. + * + * MEDIUM #1 coverage: tier_lock is per-provider (Record). + * Locking provider "agy" to "ultra" must NOT affect another provider's accounts. + * + * Strategy: + * - _mockConfig is an in-memory object; each test mutates it to set the desired + * tier_lock, then imports a fresh quota-manager instance via a cache-busting + * query string. No disk writes, no process.env.CCS_HOME races. + * - mock.module for account-manager and account-safety (already used by prior + * tests in this file) cover the full transitive dependency surface. + * - config-loader-facade is NOT mocked here — quota-manager reads config via + * loadOrCreateUnifiedConfig which in turn reads CCS_HOME from env. We + * write the minimal config.yaml once per test via writeMinimalConfig() to + * a dedicated temp dir set to CCS_HOME before each test. This gives full + * control without touching the facade mock surface. + * + * Note on isolation: Bun runs each test FILE in its own worker process, so + * process.env.CCS_HOME is NOT shared across test files. The previous 2-test + * failure ("clearing tier_lock (null)" and "no tier_lock") was caused by + * within-file state leakage: an earlier test wrote {agy:'ultra',claude:'pro'} + * to disk, and invalidateConfigCache() only cleared the facade's memoisation + * cache — loadOrCreateUnifiedConfig still read the stale disk state. The fix + * is to explicitly re-write the config file in every test that needs null-lock. + */ + +import { afterAll, afterEach, beforeEach, describe, expect, it, mock } from 'bun:test'; +import * as fs from 'fs'; +import * as os from 'os'; +import * as path from 'path'; +import { invalidateConfigCache } from '../../../src/config/config-loader-facade'; + +// ============================================================================ +// Account fixtures +// ============================================================================ + +const ULTRA_ACCOUNT = { id: 'ultra-acc-1', tier: 'ultra', email: 'ultra@example.com' }; +const PRO_ACCOUNT = { id: 'pro-acc-1', tier: 'pro', email: 'pro@example.com' }; +const PRO_ACCOUNT_2 = { id: 'pro-acc-2', tier: 'pro', email: 'pro2@example.com' }; + +/** Quota objects match calculateAgyQuotaPercent shape. */ +const HEALTHY_QUOTA = { success: true, models: [{ percentage: 80 }] }; +const EXHAUSTED_QUOTA = { success: true, models: [{ percentage: 2 }] }; + +// ============================================================================ +// Mutable mock state +// ============================================================================ + +let _mockAccounts: (typeof ULTRA_ACCOUNT)[] = []; +let _mockPausedIds: Set = new Set(); + +// ============================================================================ +// Top-level mock.module registrations (file-load time) +// ============================================================================ + +mock.module('../../../src/cliproxy/accounts/account-manager', () => ({ + PROVIDERS_WITHOUT_EMAIL: [], + getAccountsRegistryPath: () => '', + getPausedDir: () => '', + getAccountTokenPath: () => '', + extractAccountIdFromTokenFile: () => '', + deriveNoEmailProviderAccountId: () => '', + generateNickname: () => '', + validateNickname: () => true, + hasAccountNameConflict: () => false, + findAccountNameMatch: () => null, + tokenFileExists: () => false, + loadAccountsRegistry: () => ({}), + saveAccountsRegistry: () => undefined, + syncRegistryWithTokenFiles: () => undefined, + registerAccount: () => undefined, + setDefaultAccount: () => undefined, + pauseAccount: () => undefined, + resumeAccount: () => undefined, + removeAccount: () => undefined, + renameAccount: () => undefined, + touchAccount: () => undefined, + setAccountTier: () => undefined, + discoverExistingAccounts: () => [], + getProviderAccounts: () => _mockAccounts, + getDefaultAccount: () => (_mockAccounts.length > 0 ? _mockAccounts[0] : null), + getAccount: (_p: string, id: string) => _mockAccounts.find((a) => a.id === id) ?? null, + findAccountByQuery: () => null, + getActiveAccounts: () => _mockAccounts, + isAccountPaused: (_p: string, id: string) => _mockPausedIds.has(id), + getAllAccountsSummary: () => [], + bulkPauseAccounts: () => ({ succeeded: [], failed: [] }), + bulkResumeAccounts: () => ({ succeeded: [], failed: [] }), + soloAccount: async () => null, +})); + +mock.module('../../../src/cliproxy/accounts/account-safety', () => ({ + restoreExpiredQuotaPauses: () => undefined, + pauseAccountForQuotaCooldown: () => false, +})); + +// ============================================================================ +// Helpers +// ============================================================================ + +/** + * Write a minimal config.yaml in tempHome/.ccs/ with the given tier_lock. + * + * tierLock: null means no locks active. + * tierLock: Record means per-provider map, e.g. { agy: 'ultra' }. + * + * Always writes a fresh file — this is the canonical "truth" for each test. + */ +function writeMinimalConfig( + tempHome: string, + tierLock: null | Record +): void { + const ccsDir = path.join(tempHome, '.ccs'); + fs.mkdirSync(ccsDir, { recursive: true }); + + let tierLockYaml: string; + if (tierLock === null) { + tierLockYaml = 'null'; + } else { + const entries = Object.entries(tierLock) + .map(([provider, tier]) => ` ${provider}: ${tier === null ? 'null' : `"${tier}"`}`) + .join('\n'); + tierLockYaml = entries.length > 0 ? `\n${entries}` : '{}'; + } + + const yaml = ` +version: 13 +setup_completed: true +quota_management: + mode: hybrid + auto: + preflight_check: true + exhaustion_threshold: 5 + tier_priority: + - ultra + - pro + - free + cooldown_minutes: 5 + manual: + paused_accounts: [] + forced_default: null + tier_lock: ${tierLockYaml} + runtime_monitor: + enabled: true + normal_interval_seconds: 300 + critical_interval_seconds: 60 + warn_threshold: 20 + exhaustion_threshold: 5 + cooldown_minutes: 5 +`.trim(); + fs.writeFileSync(path.join(ccsDir, 'config.yaml'), yaml, 'utf-8'); +} + +// ============================================================================ +// Tests +// ============================================================================ + +describe('quota-manager findHealthyAccount — tier_lock', () => { + let tempHome = ''; + let originalCcsHome: string | undefined; + let originalCcsUnified: string | undefined; + + beforeEach(() => { + tempHome = fs.mkdtempSync(path.join(os.tmpdir(), 'ccs-tier-lock-qm-')); + originalCcsHome = process.env.CCS_HOME; + originalCcsUnified = process.env.CCS_UNIFIED_CONFIG; + process.env.CCS_HOME = tempHome; + process.env.CCS_UNIFIED_CONFIG = '1'; + + // Invalidate the shared config-loader-facade memoisation cache so that + // each test reads its own freshly-written config.yaml from disk. + invalidateConfigCache(); + + // Reset mutable mock state. + _mockAccounts = [ULTRA_ACCOUNT, PRO_ACCOUNT, PRO_ACCOUNT_2]; + _mockPausedIds = new Set(); + + // Re-register mocks in beforeEach to survive any mock.restore() calls + // from other test suites running in this Bun process. + mock.module('../../../src/cliproxy/accounts/account-manager', () => ({ + PROVIDERS_WITHOUT_EMAIL: [], + getAccountsRegistryPath: () => '', + getPausedDir: () => '', + getAccountTokenPath: () => '', + extractAccountIdFromTokenFile: () => '', + deriveNoEmailProviderAccountId: () => '', + generateNickname: () => '', + validateNickname: () => true, + hasAccountNameConflict: () => false, + findAccountNameMatch: () => null, + tokenFileExists: () => false, + loadAccountsRegistry: () => ({}), + saveAccountsRegistry: () => undefined, + syncRegistryWithTokenFiles: () => undefined, + registerAccount: () => undefined, + setDefaultAccount: () => undefined, + pauseAccount: () => undefined, + resumeAccount: () => undefined, + removeAccount: () => undefined, + renameAccount: () => undefined, + touchAccount: () => undefined, + setAccountTier: () => undefined, + discoverExistingAccounts: () => [], + getProviderAccounts: () => _mockAccounts, + getDefaultAccount: () => (_mockAccounts.length > 0 ? _mockAccounts[0] : null), + getAccount: (_p: string, id: string) => _mockAccounts.find((a) => a.id === id) ?? null, + findAccountByQuery: () => null, + getActiveAccounts: () => _mockAccounts, + isAccountPaused: (_p: string, id: string) => _mockPausedIds.has(id), + getAllAccountsSummary: () => [], + bulkPauseAccounts: () => ({ succeeded: [], failed: [] }), + bulkResumeAccounts: () => ({ succeeded: [], failed: [] }), + soloAccount: async () => null, + })); + + mock.module('../../../src/cliproxy/accounts/account-safety', () => ({ + restoreExpiredQuotaPauses: () => undefined, + pauseAccountForQuotaCooldown: () => false, + })); + }); + + afterAll(() => { + mock.restore(); + }); + + afterEach(() => { + if (originalCcsHome !== undefined) process.env.CCS_HOME = originalCcsHome; + else delete process.env.CCS_HOME; + if (originalCcsUnified !== undefined) process.env.CCS_UNIFIED_CONFIG = originalCcsUnified; + else delete process.env.CCS_UNIFIED_CONFIG; + + if (tempHome && fs.existsSync(tempHome)) { + fs.rmSync(tempHome, { recursive: true, force: true }); + } + }); + + // ------------------------------------------------------------------------- + // tier_lock = { agy: "pro" } — ultra must be excluded + // ------------------------------------------------------------------------- + + it('tier_lock="pro" — only pro accounts are candidates (ultra excluded)', async () => { + writeMinimalConfig(tempHome, { agy: 'pro' }); + _mockAccounts = [ULTRA_ACCOUNT, PRO_ACCOUNT]; + + const uid = `lock-pro-${Date.now()}-${Math.random()}`; + const { findHealthyAccount, setCachedQuota } = await import( + `../../../src/cliproxy/quota/quota-manager?${uid}` + ); + + setCachedQuota('agy', ULTRA_ACCOUNT.id, HEALTHY_QUOTA as never); + setCachedQuota('agy', PRO_ACCOUNT.id, HEALTHY_QUOTA as never); + + const result = await findHealthyAccount('agy', []); + expect(result).not.toBeNull(); + expect(result!.tier).toBe('pro'); + expect(result!.id).toBe(PRO_ACCOUNT.id); + }); + + // ------------------------------------------------------------------------- + // tier_lock specifies a tier with zero matching accounts + // ------------------------------------------------------------------------- + + it('tier_lock="ultra" with only pro accounts → returns null (no cross-tier fallback)', async () => { + writeMinimalConfig(tempHome, { agy: 'ultra' }); + _mockAccounts = [PRO_ACCOUNT, PRO_ACCOUNT_2]; + + const uid = `lock-ultra-no-ultra-${Date.now()}-${Math.random()}`; + const { findHealthyAccount, setCachedQuota } = await import( + `../../../src/cliproxy/quota/quota-manager?${uid}` + ); + + setCachedQuota('agy', PRO_ACCOUNT.id, HEALTHY_QUOTA as never); + setCachedQuota('agy', PRO_ACCOUNT_2.id, HEALTHY_QUOTA as never); + + const result = await findHealthyAccount('agy', []); + // tier_lock strictly enforced — no ultra accounts → null + expect(result).toBeNull(); + }); + + // ------------------------------------------------------------------------- + // tier_lock respects existing paused/exclude filters within locked tier + // ------------------------------------------------------------------------- + + it('tier_lock="pro" — paused pro accounts are still excluded', async () => { + writeMinimalConfig(tempHome, { agy: 'pro' }); + _mockAccounts = [PRO_ACCOUNT, PRO_ACCOUNT_2]; + _mockPausedIds = new Set([PRO_ACCOUNT.id]); + + const uid = `lock-pro-paused-${Date.now()}-${Math.random()}`; + const { findHealthyAccount, setCachedQuota } = await import( + `../../../src/cliproxy/quota/quota-manager?${uid}` + ); + + setCachedQuota('agy', PRO_ACCOUNT.id, HEALTHY_QUOTA as never); + setCachedQuota('agy', PRO_ACCOUNT_2.id, HEALTHY_QUOTA as never); + + const result = await findHealthyAccount('agy', []); + expect(result).not.toBeNull(); + expect(result!.id).toBe(PRO_ACCOUNT_2.id); + }); + + it('tier_lock="pro" — exclude list still removes accounts within the locked tier', async () => { + writeMinimalConfig(tempHome, { agy: 'pro' }); + _mockAccounts = [PRO_ACCOUNT, PRO_ACCOUNT_2]; + + const uid = `lock-pro-exclude-${Date.now()}-${Math.random()}`; + const { findHealthyAccount, setCachedQuota } = await import( + `../../../src/cliproxy/quota/quota-manager?${uid}` + ); + + setCachedQuota('agy', PRO_ACCOUNT.id, HEALTHY_QUOTA as never); + setCachedQuota('agy', PRO_ACCOUNT_2.id, HEALTHY_QUOTA as never); + + const result = await findHealthyAccount('agy', [PRO_ACCOUNT.id]); + expect(result).not.toBeNull(); + expect(result!.id).toBe(PRO_ACCOUNT_2.id); + }); + + // ------------------------------------------------------------------------- + // Clearing tier_lock restores prior behavior + // + // Key: writeMinimalConfig(tempHome, null) is called immediately before + // findHealthyAccount, AFTER setting _mockAccounts. This ensures the disk + // state is null even if a prior test in this file left stale content. + // ------------------------------------------------------------------------- + + it('clearing tier_lock (null) allows both tiers as candidates again', async () => { + _mockAccounts = [PRO_ACCOUNT]; + // Write null-lock config as the very last disk op before the import so no + // within-file test ordering can leave stale { agy: ... } on disk. + writeMinimalConfig(tempHome, null); + invalidateConfigCache(); + + const uid = `lock-cleared-${Date.now()}-${Math.random()}`; + const { findHealthyAccount, setCachedQuota } = await import( + `../../../src/cliproxy/quota/quota-manager?${uid}` + ); + + setCachedQuota('agy', PRO_ACCOUNT.id, HEALTHY_QUOTA as never); + + const result = await findHealthyAccount('agy', []); + expect(result).not.toBeNull(); + expect(result!.id).toBe(PRO_ACCOUNT.id); + }); + + // ------------------------------------------------------------------------- + // Exhausted locked-tier accounts: no cross-tier fallback + // ------------------------------------------------------------------------- + + it('tier_lock="pro" — exhausted pro + healthy ultra → returns null (no cross-tier fallback)', async () => { + writeMinimalConfig(tempHome, { agy: 'pro' }); + _mockAccounts = [ULTRA_ACCOUNT, PRO_ACCOUNT]; + + const uid = `lock-pro-exhausted-${Date.now()}-${Math.random()}`; + const { findHealthyAccount, setCachedQuota } = await import( + `../../../src/cliproxy/quota/quota-manager?${uid}` + ); + + // Pro is exhausted (2%), ultra is healthy (80%) — tier_lock must block ultra + setCachedQuota('agy', PRO_ACCOUNT.id, EXHAUSTED_QUOTA as never); + setCachedQuota('agy', ULTRA_ACCOUNT.id, HEALTHY_QUOTA as never); + + const result = await findHealthyAccount('agy', []); + // No healthy pro accounts; must NOT fall back to ultra + expect(result).toBeNull(); + }); + + // ------------------------------------------------------------------------- + // Baseline: no tier_lock — any available healthy account is returned. + // + // Same pattern: write null-lock and invalidate cache immediately before import. + // ------------------------------------------------------------------------- + + it('no tier_lock — at least one account is returned (tier filter is inactive)', async () => { + _mockAccounts = [PRO_ACCOUNT]; + // Write null-lock as the very last disk op before the import. + writeMinimalConfig(tempHome, null); + invalidateConfigCache(); + + const uid = `no-lock-${Date.now()}-${Math.random()}`; + const { findHealthyAccount, setCachedQuota } = await import( + `../../../src/cliproxy/quota/quota-manager?${uid}` + ); + + setCachedQuota('agy', PRO_ACCOUNT.id, HEALTHY_QUOTA as never); + + const result = await findHealthyAccount('agy', []); + expect(result).not.toBeNull(); + expect(result!.tier).toBe('pro'); + }); + + // ------------------------------------------------------------------------- + // MEDIUM #1: cross-provider isolation + // Locking "agy" to "ultra" must NOT filter accounts for an unlocked provider. + // ------------------------------------------------------------------------- + + it('tier_lock on agy does NOT affect another provider — pro accounts are still selectable for unlocked provider', async () => { + writeMinimalConfig(tempHome, { agy: 'ultra' }); + // Only pro accounts — agy locked to ultra means no result. + _mockAccounts = [PRO_ACCOUNT, PRO_ACCOUNT_2]; + + const uid = `cross-provider-${Date.now()}-${Math.random()}`; + const { findHealthyAccount, setCachedQuota } = await import( + `../../../src/cliproxy/quota/quota-manager?${uid}` + ); + + setCachedQuota('agy', PRO_ACCOUNT.id, HEALTHY_QUOTA as never); + setCachedQuota('agy', PRO_ACCOUNT_2.id, HEALTHY_QUOTA as never); + + // agy IS locked to ultra → no ultra accounts → null + const agyLocked = await findHealthyAccount('agy', []); + expect(agyLocked).toBeNull(); + + // Prove isolation: rewrite config with lock on a different key only. + // getTierLockForProvider(config.manual, 'agy') returns null when 'agy' is + // absent from the map — pro accounts become candidates again. + writeMinimalConfig(tempHome, { claude: 'ultra' }); // only claude locked, not agy + invalidateConfigCache(); + + // agy has no lock entry → pro accounts are candidates + const agyUnlocked = await findHealthyAccount('agy', []); + expect(agyUnlocked).not.toBeNull(); + expect(agyUnlocked!.tier).toBe('pro'); + }); +}); diff --git a/tests/unit/web-server/tier-lock-routes.test.ts b/tests/unit/web-server/tier-lock-routes.test.ts new file mode 100644 index 00000000..55736614 --- /dev/null +++ b/tests/unit/web-server/tier-lock-routes.test.ts @@ -0,0 +1,200 @@ +/** + * Phase 2: tier-lock endpoint tests + * + * POST /api/accounts/tier-lock + * body: { tier: string|null, provider?: string } + * + * Tests: + * - sets tier_lock in config and returns it + * - clears tier_lock when tier is null + * - rejects missing provider + * - rejects invalid provider + * - rejects unknown tier strings (typos must 400, not silently persist) + * - persists across config reads (config write path) as per-provider map + * - locking one provider does NOT affect another provider's lock entry + */ + +import { afterEach, beforeEach, describe, expect, it } from 'bun:test'; +import express from 'express'; +import * as fs from 'fs'; +import * as os from 'os'; +import * as path from 'path'; +import type { Server } from 'http'; + +async function postJson(baseUrl: string, routePath: string, body: unknown): Promise { + return fetch(`${baseUrl}${routePath}`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(body), + }); +} + +describe('POST /api/accounts/tier-lock', () => { + let server: Server; + let baseUrl = ''; + let tempHome = ''; + let originalCcsHome: string | undefined; + let originalCcsUnified: string | undefined; + + beforeEach(async () => { + tempHome = fs.mkdtempSync(path.join(os.tmpdir(), 'ccs-tier-lock-routes-')); + originalCcsHome = process.env.CCS_HOME; + originalCcsUnified = process.env.CCS_UNIFIED_CONFIG; + + process.env.CCS_HOME = tempHome; + process.env.CCS_UNIFIED_CONFIG = '1'; + + // No account-manager mock: the empty temp CCS_HOME yields zero CLIProxy + // accounts naturally, and the tier-lock endpoint only validates the + // provider id and writes config. Avoiding mock.module here is deliberate — + // Bun's mock.restore() does NOT unwind mock.module, so a global account + // manager mock would leak into later test files in the same process. + const { default: accountRoutes } = await import( + `../../../src/web-server/routes/account-routes?tier-lock-test=${Date.now()}-${Math.random()}` + ); + + const app = express(); + app.use(express.json()); + app.use('/api/accounts', accountRoutes); + + server = await new Promise((resolve, reject) => { + const instance = app.listen(0, '127.0.0.1'); + instance.once('error', reject); + instance.once('listening', () => resolve(instance)); + }); + + const address = server.address(); + if (!address || typeof address === 'string') { + throw new Error('Unable to resolve test server port'); + } + baseUrl = `http://127.0.0.1:${address.port}`; + }); + + afterEach(async () => { + await new Promise((resolve) => server.close(() => resolve())); + + if (originalCcsHome !== undefined) process.env.CCS_HOME = originalCcsHome; + else delete process.env.CCS_HOME; + if (originalCcsUnified !== undefined) process.env.CCS_UNIFIED_CONFIG = originalCcsUnified; + else delete process.env.CCS_UNIFIED_CONFIG; + + if (tempHome && fs.existsSync(tempHome)) { + fs.rmSync(tempHome, { recursive: true, force: true }); + } + }); + + it('sets tier_lock to a named tier and returns it', async () => { + const res = await postJson(baseUrl, '/api/accounts/tier-lock', { + provider: 'agy', + tier: 'ultra', + }); + expect(res.status).toBe(200); + const body = (await res.json()) as { provider: string; tier_lock: string | null }; + expect(body.provider).toBe('agy'); + expect(body.tier_lock).toBe('ultra'); + }); + + it('clears tier_lock when tier is null', async () => { + // Set first + await postJson(baseUrl, '/api/accounts/tier-lock', { provider: 'agy', tier: 'pro' }); + + // Then clear + const res = await postJson(baseUrl, '/api/accounts/tier-lock', { + provider: 'agy', + tier: null, + }); + expect(res.status).toBe(200); + const body = (await res.json()) as { provider: string; tier_lock: string | null }; + expect(body.provider).toBe('agy'); + expect(body.tier_lock).toBeNull(); + }); + + it('rejects missing provider with 400', async () => { + const res = await postJson(baseUrl, '/api/accounts/tier-lock', { tier: 'pro' }); + expect(res.status).toBe(400); + const body = (await res.json()) as { error: string }; + expect(body.error).toMatch(/provider/i); + }); + + it('rejects invalid provider with 400', async () => { + const res = await postJson(baseUrl, '/api/accounts/tier-lock', { + provider: 'notreal', + tier: 'pro', + }); + expect(res.status).toBe(400); + const body = (await res.json()) as { error: string }; + expect(body.error).toMatch(/provider/i); + }); + + it('rejects missing tier field (no tier key at all) with 400', async () => { + const res = await postJson(baseUrl, '/api/accounts/tier-lock', { provider: 'agy' }); + expect(res.status).toBe(400); + const body = (await res.json()) as { error: string }; + expect(body.error).toMatch(/tier/i); + }); + + it('rejects non-string non-null tier with 400', async () => { + const res = await postJson(baseUrl, '/api/accounts/tier-lock', { + provider: 'agy', + tier: 42, + }); + expect(res.status).toBe(400); + const body = (await res.json()) as { error: string }; + expect(body.error).toMatch(/tier/i); + }); + + it('rejects unknown tier string (typo) with 400', async () => { + // "Ultra" (capital U) is not a valid AccountTier — must 400, not silently persist + const res = await postJson(baseUrl, '/api/accounts/tier-lock', { + provider: 'agy', + tier: 'Ultra', + }); + expect(res.status).toBe(400); + const body = (await res.json()) as { error: string }; + expect(body.error).toMatch(/tier/i); + }); + + it('rejects "premium" (unknown tier) with 400', async () => { + const res = await postJson(baseUrl, '/api/accounts/tier-lock', { + provider: 'agy', + tier: 'premium', + }); + expect(res.status).toBe(400); + const body = (await res.json()) as { error: string }; + expect(body.error).toMatch(/tier/i); + }); + + it('persists tier_lock as per-provider map in the config', async () => { + await postJson(baseUrl, '/api/accounts/tier-lock', { provider: 'agy', tier: 'pro' }); + + // Read the config directly to confirm per-provider persistence + const { loadOrCreateUnifiedConfig } = await import( + `../../../src/config/config-loader-facade?persist-check=${Date.now()}` + ); + const config = loadOrCreateUnifiedConfig(); + const tierLock = config.quota_management?.manual?.tier_lock; + // Must be a map, not a bare string + expect(typeof tierLock).toBe('object'); + expect((tierLock as Record)['agy']).toBe('pro'); + }); + + it('tier_lock is persisted as a per-provider map entry (not a global string)', async () => { + // Lock agy to ultra — verify the map structure has only agy set + await postJson(baseUrl, '/api/accounts/tier-lock', { provider: 'agy', tier: 'ultra' }); + + const { loadOrCreateUnifiedConfig } = await import( + `../../../src/config/config-loader-facade?per-provider-map-check=${Date.now()}` + ); + const config = loadOrCreateUnifiedConfig(); + const tierLock = config.quota_management?.manual?.tier_lock; + + // Must be a map, not a bare string + expect(typeof tierLock).toBe('object'); + expect(tierLock).not.toBeNull(); + // The agy entry must be set + expect((tierLock as Record)['agy']).toBe('ultra'); + // Providers not explicitly locked must not appear in the map + expect((tierLock as Record)['codex'] ?? null).toBeNull(); + expect((tierLock as Record)['gemini'] ?? null).toBeNull(); + }); +});