mirror of
https://github.com/tiennm99/ccs.git
synced 2026-07-17 12:16:59 +00:00
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.
This commit is contained in:
@@ -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<Prefli
|
||||
}
|
||||
}
|
||||
|
||||
// When a tier is locked for this specific provider, the default account must
|
||||
// match the locked tier. If it doesn't, route to a healthy account in the
|
||||
// locked tier instead. Locks are per-provider: locking "agy" to "ultra"
|
||||
// does NOT constrain "claude", "codex", "gemini", or "ghcp".
|
||||
// Graceful degradation: if no locked-tier account is available, fall through
|
||||
// to the default (don't block the request entirely).
|
||||
const tierLock = getTierLockForProvider(quotaConfig.manual, provider);
|
||||
if (tierLock !== null && (defaultAccount.tier || 'unknown') !== tierLock) {
|
||||
const lockedTierAccount = await findHealthyAccount(provider, []);
|
||||
if (lockedTierAccount) {
|
||||
setDefaultAccount(provider, lockedTierAccount.id);
|
||||
touchAccount(provider, lockedTierAccount.id);
|
||||
return {
|
||||
proceed: true,
|
||||
accountId: lockedTierAccount.id,
|
||||
switchedFrom: defaultAccount.id,
|
||||
reason: `Tier lock: selected ${tierLock} account`,
|
||||
};
|
||||
}
|
||||
// No locked-tier account available — fall through and use default
|
||||
}
|
||||
|
||||
// Check if default is paused
|
||||
if (isAccountPaused(provider, defaultAccount.id)) {
|
||||
return await findAndSwitch(provider, defaultAccount.id, 'Default account is paused');
|
||||
|
||||
@@ -351,5 +351,24 @@ export function generateYamlWithComments(config: UnifiedConfig): string {
|
||||
lines.push('');
|
||||
}
|
||||
|
||||
// Quota management section (hybrid auto+manual account selection)
|
||||
if (config.quota_management) {
|
||||
lines.push('# ----------------------------------------------------------------------------');
|
||||
lines.push('# Quota Management: Hybrid auto+manual account selection for multi-account setups');
|
||||
lines.push('# mode: auto | manual | hybrid (default: hybrid)');
|
||||
lines.push('# manual.tier_lock: per-provider tier lock map (e.g. { agy: "ultra" })');
|
||||
lines.push('# Configure via: POST /api/accounts/tier-lock');
|
||||
lines.push('# ----------------------------------------------------------------------------');
|
||||
lines.push(
|
||||
yaml
|
||||
.dump(
|
||||
{ quota_management: config.quota_management },
|
||||
{ indent: 2, lineWidth: -1, quotingType: '"' }
|
||||
)
|
||||
.trim()
|
||||
);
|
||||
lines.push('');
|
||||
}
|
||||
|
||||
return lines.join('\n');
|
||||
}
|
||||
|
||||
@@ -52,8 +52,20 @@ export interface ManualQuotaConfig {
|
||||
paused_accounts: string[];
|
||||
/** Force use of specific account (overrides auto-selection) */
|
||||
forced_default: string | null;
|
||||
/** Lock to specific tier only */
|
||||
tier_lock: string | null;
|
||||
/**
|
||||
* Per-provider tier lock map.
|
||||
*
|
||||
* Keys are provider IDs (e.g. "agy", "claude"). Values are the tier name to
|
||||
* lock that provider to (e.g. "ultra", "pro"), or null to clear the lock for
|
||||
* that provider.
|
||||
*
|
||||
* Only providers present in the map are affected; other providers retain
|
||||
* normal tier-priority failover. A null/absent map means no locks are active.
|
||||
*
|
||||
* Legacy shape (bare string | null): treated as no lock. The old global
|
||||
* string value predates per-provider locking and is ignored on read.
|
||||
*/
|
||||
tier_lock: Record<string, string | null> | 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<string, string | null>)[provider];
|
||||
return value ?? null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Default runtime monitor configuration.
|
||||
*/
|
||||
|
||||
@@ -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<string> = new Set<AccountTier>([
|
||||
'free',
|
||||
'pro',
|
||||
'ultra',
|
||||
'unknown',
|
||||
]);
|
||||
|
||||
const router = Router();
|
||||
|
||||
@@ -633,4 +646,107 @@ router.post('/solo', async (req: Request, res: Response): Promise<void> => {
|
||||
}
|
||||
});
|
||||
|
||||
/**
|
||||
* 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<string, string | null>) };
|
||||
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;
|
||||
|
||||
@@ -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<provider, tier|null>).
|
||||
* 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<string> = 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<string, string | null>
|
||||
): 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');
|
||||
});
|
||||
});
|
||||
@@ -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<Response> {
|
||||
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<Server>((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<void>((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<string, string | null>)['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<string, string | null>)['agy']).toBe('ultra');
|
||||
// Providers not explicitly locked must not appear in the map
|
||||
expect((tierLock as Record<string, string | null>)['codex'] ?? null).toBeNull();
|
||||
expect((tierLock as Record<string, string | null>)['gemini'] ?? null).toBeNull();
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user