mirror of
https://github.com/tiennm99/ccs.git
synced 2026-07-16 02:11:28 +00:00
feat(cliproxy): runtime quota monitoring during active sessions (#529)
* feat(cliproxy): add runtime quota monitoring during active sessions Adds adaptive background quota polling to detect and respond to quota exhaustion during active CLIProxy sessions. Prevents rate-limit-driven account bans by auto-cooling exhausted accounts and switching defaults. - Adaptive polling: 300s normal, 60s at 20% threshold, stops at 0% - Stderr warnings at 20%, boxed exhaustion alerts at 0% - Cooldown + default switch on exhaustion (existing patterns) - Configurable via quota_management.runtime_monitor in config.yaml - Timer.unref() prevents blocking process exit - monitorStopped guard for in-flight poll safety Closes #524 * fix: address code review feedback (attempt 1/5) - M1: Round quotaPercent display with Math.round() to avoid ugly floats - M2: Rename exhaust_threshold -> exhaustion_threshold for consistency with existing auto.exhaustion_threshold config field - M3: Replace async not.toThrow() with direct await assertion pattern * fix: address code review feedback (attempt 2/5) - Remove .claude/agent-memory/ from tracking and add to .gitignore - Unify cooldown_minutes default to 5 (was 10 in runtime_monitor, 5 in auto) - Add threshold validation in startQuotaMonitor (warn > exhaustion) - Document intentional post-switch monitoring gap in code comment
This commit is contained in:
@@ -33,6 +33,7 @@ pnpm-lock.yaml
|
||||
package-lock.json
|
||||
|
||||
.claude/active-plan
|
||||
.claude/agent-memory/
|
||||
|
||||
# Logs directory
|
||||
logs/
|
||||
|
||||
@@ -374,3 +374,97 @@ export function maskEmail(email: string): string {
|
||||
function truncate(str: string, maxLen: number): string {
|
||||
return str.length > maxLen ? str.slice(0, maxLen - 3) + '...' : str;
|
||||
}
|
||||
|
||||
// --- Quota Exhaustion Handling ---
|
||||
|
||||
/**
|
||||
* Write boxed quota warning to stderr (20% threshold).
|
||||
* Uses process.stderr.write() to work alongside inherited stdio.
|
||||
* ASCII-only output (no emojis) per project constraints.
|
||||
*/
|
||||
export function writeQuotaWarning(accountId: string, quotaPercent: number): void {
|
||||
const masked = maskEmail(accountId);
|
||||
const lines = [
|
||||
`[!] Quota Low: ${masked} (${Math.round(quotaPercent)}% remaining)`,
|
||||
` Next session will use a different account if available`,
|
||||
];
|
||||
const maxLen = Math.max(...lines.map((l) => l.length));
|
||||
const border = '\u2550'.repeat(maxLen + 2);
|
||||
|
||||
process.stderr.write('\n');
|
||||
process.stderr.write(`\u2554${border}\u2557\n`);
|
||||
for (const line of lines) {
|
||||
process.stderr.write(`\u2551 ${line.padEnd(maxLen)} \u2551\n`);
|
||||
}
|
||||
process.stderr.write(`\u255A${border}\u255D\n`);
|
||||
process.stderr.write('\n');
|
||||
}
|
||||
|
||||
/**
|
||||
* Write boxed quota exhaustion alert to stderr.
|
||||
* Called when quota falls below exhaustion_threshold — account will be cooled down.
|
||||
*/
|
||||
function writeQuotaExhausted(
|
||||
accountId: string,
|
||||
switchedTo: string | null,
|
||||
cooldownMinutes: number
|
||||
): void {
|
||||
const masked = maskEmail(accountId);
|
||||
const lines = [`[X] Quota Exhausted: ${masked}`, ` Cooldown: ${cooldownMinutes} minutes`];
|
||||
if (switchedTo) {
|
||||
lines.push(` Next session default: ${maskEmail(switchedTo)}`);
|
||||
} else {
|
||||
lines.push(` No alternative accounts available`);
|
||||
}
|
||||
|
||||
const maxLen = Math.max(...lines.map((l) => l.length));
|
||||
const border = '\u2550'.repeat(maxLen + 2);
|
||||
|
||||
process.stderr.write('\n');
|
||||
process.stderr.write(`\u2554${border}\u2557\n`);
|
||||
for (const line of lines) {
|
||||
process.stderr.write(`\u2551 ${line.padEnd(maxLen)} \u2551\n`);
|
||||
}
|
||||
process.stderr.write(`\u255A${border}\u255D\n`);
|
||||
process.stderr.write('\n');
|
||||
}
|
||||
|
||||
/**
|
||||
* Handle quota exhaustion for an active session.
|
||||
* Applies cooldown to exhausted account, finds healthy alternative,
|
||||
* switches default, and alerts user via stderr.
|
||||
*
|
||||
* @returns switchedTo account ID or null if no alternatives
|
||||
*/
|
||||
export async function handleQuotaExhaustion(
|
||||
provider: CLIProxyProvider,
|
||||
accountId: string,
|
||||
cooldownMinutes: number
|
||||
): Promise<{ switchedTo: string | null; reason: string }> {
|
||||
// Dynamic imports to avoid circular dependencies
|
||||
const { applyCooldown, findHealthyAccount } = await import('./quota-manager');
|
||||
const { setDefaultAccount, touchAccount } = await import('./account-manager');
|
||||
|
||||
// Apply cooldown to exhausted account
|
||||
applyCooldown(provider, accountId, cooldownMinutes);
|
||||
|
||||
// Find healthy alternative
|
||||
const alternative = await findHealthyAccount(provider, [accountId]);
|
||||
|
||||
if (alternative) {
|
||||
setDefaultAccount(provider, alternative.id);
|
||||
touchAccount(provider, alternative.id);
|
||||
writeQuotaExhausted(accountId, alternative.id, cooldownMinutes);
|
||||
return {
|
||||
switchedTo: alternative.id,
|
||||
reason: `Quota exhausted, switched to ${alternative.id}`,
|
||||
};
|
||||
}
|
||||
|
||||
// No alternatives — warn but continue (graceful degradation)
|
||||
writeQuotaExhausted(accountId, null, cooldownMinutes);
|
||||
return {
|
||||
switchedTo: null,
|
||||
reason: 'Quota exhausted, no alternatives available',
|
||||
};
|
||||
}
|
||||
|
||||
@@ -793,6 +793,15 @@ export async function execClaudeWithCLIProxy(
|
||||
});
|
||||
}
|
||||
|
||||
// 12b. Start runtime quota monitor (adaptive polling during session)
|
||||
if (!skipLocalAuth) {
|
||||
const { startQuotaMonitor } = await import('../quota-manager');
|
||||
const monitorAccount = getDefaultAccount(provider);
|
||||
if (monitorAccount) {
|
||||
startQuotaMonitor(provider, monitorAccount.id);
|
||||
}
|
||||
}
|
||||
|
||||
// 13. Setup cleanup handlers
|
||||
setupCleanupHandlers(
|
||||
claude,
|
||||
|
||||
@@ -21,6 +21,7 @@ import {
|
||||
import { detectRunningProxy, waitForProxyHealthy, reclaimOrphanedProxy } from '../proxy-detector';
|
||||
import { withStartupLock } from '../startup-lock';
|
||||
import { killProcessOnPort } from '../../utils/platform-commands';
|
||||
import { stopQuotaMonitor } from '../quota-manager';
|
||||
|
||||
export interface ProxySessionResult {
|
||||
sessionId?: string;
|
||||
@@ -184,6 +185,7 @@ export function setupCleanupHandlers(
|
||||
};
|
||||
|
||||
const cleanup = () => {
|
||||
stopQuotaMonitor();
|
||||
log('Parent signal received, cleaning up');
|
||||
|
||||
if (
|
||||
@@ -214,6 +216,7 @@ export function setupCleanupHandlers(
|
||||
};
|
||||
|
||||
claude.on('exit', (code, signal) => {
|
||||
stopQuotaMonitor();
|
||||
log(`Claude exited: code=${code}, signal=${signal}`);
|
||||
|
||||
if (
|
||||
@@ -250,6 +253,7 @@ export function setupCleanupHandlers(
|
||||
});
|
||||
|
||||
claude.on('error', (error) => {
|
||||
stopQuotaMonitor();
|
||||
console.error(require('../../utils/ui').fail(`Claude CLI error: ${error}`));
|
||||
|
||||
if (
|
||||
|
||||
@@ -23,6 +23,7 @@ import {
|
||||
type AccountInfo,
|
||||
} from './account-manager';
|
||||
import { loadOrCreateUnifiedConfig } from '../config/unified-config-loader';
|
||||
import type { RuntimeMonitorConfig } from '../config/unified-config-types';
|
||||
|
||||
// ============================================================================
|
||||
// QUOTA CACHE (30-second TTL)
|
||||
@@ -416,3 +417,137 @@ export async function getQuotaStatus(provider: CLIProxyProvider): Promise<{
|
||||
|
||||
return { accounts: results };
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// RUNTIME QUOTA MONITOR (adaptive polling during active sessions)
|
||||
// ============================================================================
|
||||
|
||||
/** Active monitor timer (null = not running) */
|
||||
let monitorTimer: ReturnType<typeof setTimeout> | null = null;
|
||||
|
||||
/** Tracks if warning was shown this session (avoid spam) */
|
||||
let hasWarnedThisSession = false;
|
||||
|
||||
/** Guards against in-flight poll callbacks running after stop */
|
||||
let monitorStopped = false;
|
||||
|
||||
/**
|
||||
* Schedule next quota poll with adaptive interval.
|
||||
* Uses setTimeout chain (not setInterval) for dynamic interval switching.
|
||||
*/
|
||||
function scheduleNextPoll(
|
||||
provider: CLIProxyProvider,
|
||||
accountId: string,
|
||||
monitorConfig: RuntimeMonitorConfig,
|
||||
intervalMs: number
|
||||
): void {
|
||||
monitorTimer = setTimeout(async () => {
|
||||
// Guard: skip if monitor was stopped while this callback was queued
|
||||
if (monitorStopped) return;
|
||||
|
||||
try {
|
||||
const quota = await fetchQuotaWithDedup(provider, accountId);
|
||||
if (monitorStopped) return; // Re-check after async fetch
|
||||
const avgQuota = calculateAverageQuota(quota) ?? 100;
|
||||
|
||||
if (avgQuota <= monitorConfig.exhaustion_threshold) {
|
||||
// EXHAUSTED: cooldown + switch default + stop monitoring.
|
||||
// NOTE: Monitor stops here intentionally. The current session continues
|
||||
// on the exhausted account (can't hot-swap mid-session). The switched
|
||||
// default only takes effect on next session start via preflightCheck().
|
||||
const { handleQuotaExhaustion } = await import('./account-safety');
|
||||
await handleQuotaExhaustion(provider, accountId, monitorConfig.cooldown_minutes);
|
||||
monitorTimer = null;
|
||||
return; // Stop polling
|
||||
}
|
||||
|
||||
if (avgQuota <= monitorConfig.warn_threshold) {
|
||||
// WARNING: switch to critical interval, warn once
|
||||
if (!hasWarnedThisSession) {
|
||||
const { writeQuotaWarning } = await import('./account-safety');
|
||||
writeQuotaWarning(accountId, avgQuota);
|
||||
hasWarnedThisSession = true;
|
||||
}
|
||||
scheduleNextPoll(
|
||||
provider,
|
||||
accountId,
|
||||
monitorConfig,
|
||||
monitorConfig.critical_interval_seconds * 1000
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
// HEALTHY: keep normal interval
|
||||
scheduleNextPoll(
|
||||
provider,
|
||||
accountId,
|
||||
monitorConfig,
|
||||
monitorConfig.normal_interval_seconds * 1000
|
||||
);
|
||||
} catch {
|
||||
// API failure: silently reschedule at same interval
|
||||
scheduleNextPoll(provider, accountId, monitorConfig, intervalMs);
|
||||
}
|
||||
}, intervalMs);
|
||||
|
||||
// Prevent monitor from keeping Node.js process alive
|
||||
if (monitorTimer && typeof monitorTimer === 'object' && 'unref' in monitorTimer) {
|
||||
monitorTimer.unref();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Start adaptive quota monitor for an active session.
|
||||
* Polls at normal_interval (300s) when healthy, switches to
|
||||
* critical_interval (60s) when quota hits warn_threshold (20%).
|
||||
* Auto-stops on exhaustion or when stopQuotaMonitor() is called.
|
||||
*
|
||||
* Only monitors 'agy' provider (only one with quota API).
|
||||
* No-op for other providers, manual mode, or if disabled in config.
|
||||
*/
|
||||
export function startQuotaMonitor(provider: CLIProxyProvider, accountId: string): void {
|
||||
// Only Antigravity supports quota
|
||||
if (provider !== 'agy') return;
|
||||
|
||||
// Prevent duplicate monitors
|
||||
if (monitorTimer) return;
|
||||
|
||||
const config = loadOrCreateUnifiedConfig();
|
||||
const quotaConfig = config.quota_management;
|
||||
|
||||
// Skip if config missing (shouldn't happen with defaults)
|
||||
if (!quotaConfig) return;
|
||||
|
||||
// Skip if manual mode or runtime monitor disabled
|
||||
if (quotaConfig.mode === 'manual') return;
|
||||
if (!quotaConfig.runtime_monitor?.enabled) return;
|
||||
|
||||
// Validate thresholds: warn must be > exhaustion to avoid immediate exhaustion on warning
|
||||
const monitorConfig = quotaConfig.runtime_monitor;
|
||||
if (monitorConfig.warn_threshold <= monitorConfig.exhaustion_threshold) {
|
||||
return; // Invalid config — skip monitoring silently (logged at config level)
|
||||
}
|
||||
|
||||
hasWarnedThisSession = false;
|
||||
monitorStopped = false;
|
||||
|
||||
// Start first poll at normal interval
|
||||
scheduleNextPoll(
|
||||
provider,
|
||||
accountId,
|
||||
quotaConfig.runtime_monitor,
|
||||
quotaConfig.runtime_monitor.normal_interval_seconds * 1000
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Stop the runtime quota monitor. Safe to call multiple times.
|
||||
*/
|
||||
export function stopQuotaMonitor(): void {
|
||||
monitorStopped = true;
|
||||
if (monitorTimer) {
|
||||
clearTimeout(monitorTimer);
|
||||
monitorTimer = null;
|
||||
}
|
||||
hasWarnedThisSession = false;
|
||||
}
|
||||
|
||||
@@ -341,6 +341,26 @@ function mergeWithDefaults(partial: Partial<UnifiedConfig>): UnifiedConfig {
|
||||
partial.quota_management?.manual?.tier_lock ??
|
||||
DEFAULT_QUOTA_MANAGEMENT_CONFIG.manual.tier_lock,
|
||||
},
|
||||
runtime_monitor: {
|
||||
enabled:
|
||||
partial.quota_management?.runtime_monitor?.enabled ??
|
||||
DEFAULT_QUOTA_MANAGEMENT_CONFIG.runtime_monitor.enabled,
|
||||
normal_interval_seconds:
|
||||
partial.quota_management?.runtime_monitor?.normal_interval_seconds ??
|
||||
DEFAULT_QUOTA_MANAGEMENT_CONFIG.runtime_monitor.normal_interval_seconds,
|
||||
critical_interval_seconds:
|
||||
partial.quota_management?.runtime_monitor?.critical_interval_seconds ??
|
||||
DEFAULT_QUOTA_MANAGEMENT_CONFIG.runtime_monitor.critical_interval_seconds,
|
||||
warn_threshold:
|
||||
partial.quota_management?.runtime_monitor?.warn_threshold ??
|
||||
DEFAULT_QUOTA_MANAGEMENT_CONFIG.runtime_monitor.warn_threshold,
|
||||
exhaustion_threshold:
|
||||
partial.quota_management?.runtime_monitor?.exhaustion_threshold ??
|
||||
DEFAULT_QUOTA_MANAGEMENT_CONFIG.runtime_monitor.exhaustion_threshold,
|
||||
cooldown_minutes:
|
||||
partial.quota_management?.runtime_monitor?.cooldown_minutes ??
|
||||
DEFAULT_QUOTA_MANAGEMENT_CONFIG.runtime_monitor.cooldown_minutes,
|
||||
},
|
||||
},
|
||||
// Thinking config - auto/manual/off control for reasoning budget
|
||||
thinking: {
|
||||
|
||||
@@ -369,6 +369,25 @@ export interface AutoQuotaConfig {
|
||||
cooldown_minutes: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Runtime quota monitor configuration.
|
||||
* Controls adaptive polling during active sessions.
|
||||
*/
|
||||
export interface RuntimeMonitorConfig {
|
||||
/** Enable runtime monitoring during sessions (default: true) */
|
||||
enabled: boolean;
|
||||
/** Poll interval in seconds when quota > warn_threshold (default: 300) */
|
||||
normal_interval_seconds: number;
|
||||
/** Poll interval in seconds when quota <= warn_threshold (default: 60) */
|
||||
critical_interval_seconds: number;
|
||||
/** Quota percentage that triggers fast polling + warning (default: 20) */
|
||||
warn_threshold: number;
|
||||
/** Quota percentage that triggers cooldown + switch (default: 5) */
|
||||
exhaustion_threshold: number;
|
||||
/** Minutes to cooldown exhausted account (default: 10) */
|
||||
cooldown_minutes: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Manual quota management configuration.
|
||||
* User-controlled overrides for account selection.
|
||||
@@ -401,6 +420,8 @@ export interface QuotaManagementConfig {
|
||||
auto: AutoQuotaConfig;
|
||||
/** Manual mode settings */
|
||||
manual: ManualQuotaConfig;
|
||||
/** Runtime monitor settings */
|
||||
runtime_monitor: RuntimeMonitorConfig;
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -422,6 +443,18 @@ export const DEFAULT_MANUAL_QUOTA_CONFIG: ManualQuotaConfig = {
|
||||
tier_lock: null,
|
||||
};
|
||||
|
||||
/**
|
||||
* Default runtime monitor configuration.
|
||||
*/
|
||||
export const DEFAULT_RUNTIME_MONITOR_CONFIG: RuntimeMonitorConfig = {
|
||||
enabled: true,
|
||||
normal_interval_seconds: 300,
|
||||
critical_interval_seconds: 60,
|
||||
warn_threshold: 20,
|
||||
exhaustion_threshold: 5,
|
||||
cooldown_minutes: 5,
|
||||
};
|
||||
|
||||
/**
|
||||
* Default quota management configuration.
|
||||
*/
|
||||
@@ -429,6 +462,7 @@ export const DEFAULT_QUOTA_MANAGEMENT_CONFIG: QuotaManagementConfig = {
|
||||
mode: 'hybrid',
|
||||
auto: { ...DEFAULT_AUTO_QUOTA_CONFIG },
|
||||
manual: { ...DEFAULT_MANUAL_QUOTA_CONFIG },
|
||||
runtime_monitor: { ...DEFAULT_RUNTIME_MONITOR_CONFIG },
|
||||
};
|
||||
|
||||
// ============================================================================
|
||||
|
||||
@@ -0,0 +1,312 @@
|
||||
/**
|
||||
* Account Safety Quota Exhaustion Handler Tests
|
||||
*
|
||||
* Tests for handleQuotaExhaustion() and writeQuotaWarning():
|
||||
* - Cooldown application
|
||||
* - Account switching
|
||||
* - Fallback when no alternatives
|
||||
* - Warning output formatting
|
||||
* - Email masking
|
||||
*/
|
||||
|
||||
import { describe, it, expect, beforeEach, afterEach } from 'bun:test';
|
||||
import * as fs from 'fs';
|
||||
import * as os from 'os';
|
||||
import * as path from 'path';
|
||||
import { handleQuotaExhaustion, writeQuotaWarning, maskEmail } from '../../../src/cliproxy/account-safety';
|
||||
|
||||
// Setup test isolation
|
||||
let tmpDir: string;
|
||||
let origCcsHome: string | undefined;
|
||||
|
||||
beforeEach(() => {
|
||||
tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'ccs-test-exhaust-'));
|
||||
origCcsHome = process.env.CCS_HOME;
|
||||
process.env.CCS_HOME = tmpDir;
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
if (origCcsHome !== undefined) {
|
||||
process.env.CCS_HOME = origCcsHome;
|
||||
} else {
|
||||
delete process.env.CCS_HOME;
|
||||
}
|
||||
fs.rmSync(tmpDir, { recursive: true, force: true });
|
||||
});
|
||||
|
||||
// Helper: write accounts registry
|
||||
function writeRegistry(providers: Record<string, unknown>): void {
|
||||
const registryDir = path.join(tmpDir, '.ccs', 'cliproxy');
|
||||
fs.mkdirSync(registryDir, { recursive: true });
|
||||
fs.writeFileSync(
|
||||
path.join(registryDir, 'accounts.json'),
|
||||
JSON.stringify({ version: 1, providers }, null, 2)
|
||||
);
|
||||
}
|
||||
|
||||
// Helper: write unified config
|
||||
function writeConfig(quotaConfig: unknown): void {
|
||||
const configDir = path.join(tmpDir, '.ccs', 'config');
|
||||
fs.mkdirSync(configDir, { recursive: true });
|
||||
fs.writeFileSync(
|
||||
path.join(configDir, 'unified-config.json'),
|
||||
JSON.stringify({
|
||||
version: 2,
|
||||
quota_management: quotaConfig,
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
describe('Quota Exhaustion Handlers', () => {
|
||||
describe('writeQuotaWarning', () => {
|
||||
it('should write to stderr with box format', async () => {
|
||||
const stderrWrites: string[] = [];
|
||||
const originalWrite = process.stderr.write;
|
||||
process.stderr.write = ((chunk: string) => {
|
||||
stderrWrites.push(chunk);
|
||||
return true;
|
||||
}) as any;
|
||||
|
||||
writeQuotaWarning('test@gmail.com', 20);
|
||||
|
||||
process.stderr.write = originalWrite;
|
||||
|
||||
// Verify output contains account
|
||||
const fullOutput = stderrWrites.join('');
|
||||
expect(fullOutput).toContain('tes');
|
||||
expect(fullOutput).toContain('20%');
|
||||
|
||||
// Verify box borders present
|
||||
expect(fullOutput).toContain('\u2554'); // Top-left corner
|
||||
expect(fullOutput).toContain('\u2557'); // Top-right corner
|
||||
expect(fullOutput).toContain('\u255A'); // Bottom-left corner
|
||||
expect(fullOutput).toContain('\u255D'); // Bottom-right corner
|
||||
expect(fullOutput).toContain('\u2551'); // Vertical bar
|
||||
});
|
||||
|
||||
it('should mask email showing only first 3 chars', async () => {
|
||||
const stderrWrites: string[] = [];
|
||||
const originalWrite = process.stderr.write;
|
||||
process.stderr.write = ((chunk: string) => {
|
||||
stderrWrites.push(chunk);
|
||||
return true;
|
||||
}) as any;
|
||||
|
||||
writeQuotaWarning('verylongemail@example.com', 15);
|
||||
|
||||
process.stderr.write = originalWrite;
|
||||
|
||||
const fullOutput = stderrWrites.join('');
|
||||
// Should show "ver***@example.com"
|
||||
expect(fullOutput).toContain('ver***@example.com');
|
||||
expect(fullOutput).not.toContain('verylongemail@example.com');
|
||||
});
|
||||
|
||||
it('should include threshold percentage', async () => {
|
||||
const stderrWrites: string[] = [];
|
||||
const originalWrite = process.stderr.write;
|
||||
process.stderr.write = ((chunk: string) => {
|
||||
stderrWrites.push(chunk);
|
||||
return true;
|
||||
}) as any;
|
||||
|
||||
writeQuotaWarning('test@gmail.com', 5);
|
||||
|
||||
process.stderr.write = originalWrite;
|
||||
|
||||
const fullOutput = stderrWrites.join('');
|
||||
expect(fullOutput).toContain('5%');
|
||||
});
|
||||
});
|
||||
|
||||
describe('maskEmail', () => {
|
||||
it('should mask standard email', () => {
|
||||
const result = maskEmail('user@example.com');
|
||||
expect(result).toBe('use***@example.com');
|
||||
});
|
||||
|
||||
it('should handle short local part', () => {
|
||||
const result = maskEmail('ab@example.com');
|
||||
expect(result).toBe('ab***@example.com');
|
||||
});
|
||||
|
||||
it('should handle single char local part', () => {
|
||||
const result = maskEmail('a@example.com');
|
||||
expect(result).toBe('a***@example.com');
|
||||
});
|
||||
|
||||
it('should return input if no @ sign', () => {
|
||||
const result = maskEmail('not-an-email');
|
||||
expect(result).toBe('not-an-email');
|
||||
});
|
||||
|
||||
it('should return input if empty string', () => {
|
||||
const result = maskEmail('');
|
||||
expect(result).toBe('');
|
||||
});
|
||||
});
|
||||
|
||||
describe('handleQuotaExhaustion', () => {
|
||||
it('should apply cooldown to exhausted account', async () => {
|
||||
writeRegistry({
|
||||
agy: {
|
||||
default: 'exhausted@gmail.com',
|
||||
accounts: {
|
||||
'exhausted@gmail.com': {
|
||||
email: 'exhausted@gmail.com',
|
||||
tokenFile: 'agy-exhausted.json',
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
writeConfig({
|
||||
mode: 'auto',
|
||||
auto: {
|
||||
tier_priority: ['ultra', 'pro'],
|
||||
exhaustion_threshold: 5,
|
||||
cooldown_minutes: 10,
|
||||
preflight_check: true,
|
||||
},
|
||||
runtime_monitor: {
|
||||
enabled: true,
|
||||
normal_interval_seconds: 300,
|
||||
critical_interval_seconds: 60,
|
||||
warn_threshold: 20,
|
||||
exhaustion_threshold: 0,
|
||||
cooldown_minutes: 10,
|
||||
},
|
||||
});
|
||||
|
||||
const { isOnCooldown } = await import('../../../src/cliproxy/quota-manager');
|
||||
|
||||
const result = await handleQuotaExhaustion('agy', 'exhausted@gmail.com', 10);
|
||||
|
||||
// Verify cooldown was applied (account now on cooldown)
|
||||
expect(isOnCooldown('agy', 'exhausted@gmail.com')).toBe(true);
|
||||
// Should return a result with reason
|
||||
expect(result.reason).toBeDefined();
|
||||
});
|
||||
|
||||
it('should handle no alternatives gracefully', async () => {
|
||||
writeRegistry({
|
||||
agy: {
|
||||
default: 'only@gmail.com',
|
||||
accounts: {
|
||||
'only@gmail.com': {
|
||||
email: 'only@gmail.com',
|
||||
tokenFile: 'agy-only.json',
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
writeConfig({
|
||||
mode: 'auto',
|
||||
auto: {
|
||||
tier_priority: ['ultra', 'pro'],
|
||||
exhaustion_threshold: 5,
|
||||
cooldown_minutes: 10,
|
||||
preflight_check: true,
|
||||
},
|
||||
runtime_monitor: {
|
||||
enabled: true,
|
||||
normal_interval_seconds: 300,
|
||||
critical_interval_seconds: 60,
|
||||
warn_threshold: 20,
|
||||
exhaustion_threshold: 0,
|
||||
cooldown_minutes: 10,
|
||||
},
|
||||
});
|
||||
|
||||
const result = await handleQuotaExhaustion('agy', 'only@gmail.com', 10);
|
||||
|
||||
// Should return gracefully with null switched
|
||||
expect(result.switchedTo).toBeNull();
|
||||
expect(result.reason).toContain('no alternatives');
|
||||
});
|
||||
|
||||
it('should write warning to stderr', async () => {
|
||||
writeRegistry({
|
||||
agy: {
|
||||
default: 'exhausted@gmail.com',
|
||||
accounts: {
|
||||
'exhausted@gmail.com': {
|
||||
email: 'exhausted@gmail.com',
|
||||
tokenFile: 'agy-exhausted.json',
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
writeConfig({
|
||||
mode: 'auto',
|
||||
auto: {
|
||||
tier_priority: ['ultra', 'pro'],
|
||||
exhaustion_threshold: 5,
|
||||
cooldown_minutes: 10,
|
||||
preflight_check: true,
|
||||
},
|
||||
runtime_monitor: {
|
||||
enabled: true,
|
||||
normal_interval_seconds: 300,
|
||||
critical_interval_seconds: 60,
|
||||
warn_threshold: 20,
|
||||
exhaustion_threshold: 0,
|
||||
cooldown_minutes: 10,
|
||||
},
|
||||
});
|
||||
|
||||
const stderrWrites: string[] = [];
|
||||
const originalWrite = process.stderr.write;
|
||||
process.stderr.write = ((chunk: string) => {
|
||||
stderrWrites.push(chunk);
|
||||
return true;
|
||||
}) as any;
|
||||
|
||||
await handleQuotaExhaustion('agy', 'exhausted@gmail.com', 10);
|
||||
|
||||
process.stderr.write = originalWrite;
|
||||
|
||||
const fullOutput = stderrWrites.join('');
|
||||
// Should contain exhaustion indicator
|
||||
expect(fullOutput).toContain('[X]');
|
||||
});
|
||||
|
||||
it('should complete without throwing', async () => {
|
||||
writeRegistry({
|
||||
agy: {
|
||||
default: 'test@gmail.com',
|
||||
accounts: {
|
||||
'test@gmail.com': {
|
||||
email: 'test@gmail.com',
|
||||
tokenFile: 'agy-test.json',
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
writeConfig({
|
||||
mode: 'auto',
|
||||
auto: {
|
||||
tier_priority: ['ultra', 'pro'],
|
||||
exhaustion_threshold: 5,
|
||||
cooldown_minutes: 5,
|
||||
preflight_check: true,
|
||||
},
|
||||
runtime_monitor: {
|
||||
enabled: true,
|
||||
normal_interval_seconds: 300,
|
||||
critical_interval_seconds: 60,
|
||||
warn_threshold: 20,
|
||||
exhaustion_threshold: 0,
|
||||
cooldown_minutes: 5,
|
||||
},
|
||||
});
|
||||
|
||||
const result = await handleQuotaExhaustion('agy', 'test@gmail.com', 5);
|
||||
expect(result).toBeDefined();
|
||||
expect(result.switchedTo).toBeNull();
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,179 @@
|
||||
/**
|
||||
* Runtime Quota Monitor Unit Tests
|
||||
*
|
||||
* Tests the quota monitor lifecycle:
|
||||
* - startQuotaMonitor / stopQuotaMonitor behavior
|
||||
* - No-op conditions for non-agy, manual mode, disabled config
|
||||
* - Idempotent stopQuotaMonitor
|
||||
*/
|
||||
|
||||
import { describe, it, expect, beforeEach, afterEach } from 'bun:test';
|
||||
import * as fs from 'fs';
|
||||
import * as os from 'os';
|
||||
import * as path from 'path';
|
||||
import { startQuotaMonitor, stopQuotaMonitor, clearQuotaCache } from '../../../src/cliproxy/quota-manager';
|
||||
|
||||
// Setup test isolation
|
||||
let tmpDir: string;
|
||||
let origCcsHome: string | undefined;
|
||||
|
||||
beforeEach(() => {
|
||||
tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'ccs-test-monitor-'));
|
||||
origCcsHome = process.env.CCS_HOME;
|
||||
process.env.CCS_HOME = tmpDir;
|
||||
clearQuotaCache(); // Clean cache between tests
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
stopQuotaMonitor(); // Clean up any active timers
|
||||
clearQuotaCache();
|
||||
if (origCcsHome !== undefined) {
|
||||
process.env.CCS_HOME = origCcsHome;
|
||||
} else {
|
||||
delete process.env.CCS_HOME;
|
||||
}
|
||||
fs.rmSync(tmpDir, { recursive: true, force: true });
|
||||
});
|
||||
|
||||
describe('Runtime Quota Monitor', () => {
|
||||
describe('startQuotaMonitor', () => {
|
||||
it('should accept non-agy provider without throwing', () => {
|
||||
// Non-agy providers should be silently ignored
|
||||
expect(() => {
|
||||
startQuotaMonitor('gemini', 'test@gmail.com');
|
||||
}).not.toThrow();
|
||||
});
|
||||
|
||||
it('should accept agy provider without throwing', () => {
|
||||
// Setup config
|
||||
const configDir = path.join(tmpDir, '.ccs', 'config');
|
||||
fs.mkdirSync(configDir, { recursive: true });
|
||||
fs.writeFileSync(
|
||||
path.join(configDir, 'unified-config.json'),
|
||||
JSON.stringify({
|
||||
version: 2,
|
||||
quota_management: {
|
||||
mode: 'auto',
|
||||
runtime_monitor: {
|
||||
enabled: false, // Disabled to avoid actual polling
|
||||
normal_interval_seconds: 300,
|
||||
critical_interval_seconds: 60,
|
||||
warn_threshold: 20,
|
||||
exhaustion_threshold: 0,
|
||||
cooldown_minutes: 5,
|
||||
},
|
||||
},
|
||||
})
|
||||
);
|
||||
|
||||
expect(() => {
|
||||
startQuotaMonitor('agy', 'test@gmail.com');
|
||||
}).not.toThrow();
|
||||
});
|
||||
|
||||
it('should be no-op when config missing or no quota_management', () => {
|
||||
// No config file — should not throw
|
||||
expect(() => {
|
||||
startQuotaMonitor('agy', 'test@gmail.com');
|
||||
}).not.toThrow();
|
||||
});
|
||||
|
||||
it('should handle manual mode gracefully', () => {
|
||||
const configDir = path.join(tmpDir, '.ccs', 'config');
|
||||
fs.mkdirSync(configDir, { recursive: true });
|
||||
fs.writeFileSync(
|
||||
path.join(configDir, 'unified-config.json'),
|
||||
JSON.stringify({
|
||||
version: 2,
|
||||
quota_management: {
|
||||
mode: 'manual',
|
||||
runtime_monitor: {
|
||||
enabled: true,
|
||||
normal_interval_seconds: 300,
|
||||
critical_interval_seconds: 60,
|
||||
warn_threshold: 20,
|
||||
exhaustion_threshold: 0,
|
||||
cooldown_minutes: 5,
|
||||
},
|
||||
},
|
||||
})
|
||||
);
|
||||
|
||||
expect(() => {
|
||||
startQuotaMonitor('agy', 'test@gmail.com');
|
||||
}).not.toThrow();
|
||||
});
|
||||
|
||||
it('should handle disabled monitor gracefully', () => {
|
||||
const configDir = path.join(tmpDir, '.ccs', 'config');
|
||||
fs.mkdirSync(configDir, { recursive: true });
|
||||
fs.writeFileSync(
|
||||
path.join(configDir, 'unified-config.json'),
|
||||
JSON.stringify({
|
||||
version: 2,
|
||||
quota_management: {
|
||||
mode: 'auto',
|
||||
runtime_monitor: {
|
||||
enabled: false,
|
||||
normal_interval_seconds: 300,
|
||||
critical_interval_seconds: 60,
|
||||
warn_threshold: 20,
|
||||
exhaustion_threshold: 0,
|
||||
cooldown_minutes: 5,
|
||||
},
|
||||
},
|
||||
})
|
||||
);
|
||||
|
||||
expect(() => {
|
||||
startQuotaMonitor('agy', 'test@gmail.com');
|
||||
}).not.toThrow();
|
||||
});
|
||||
});
|
||||
|
||||
describe('stopQuotaMonitor', () => {
|
||||
it('should be idempotent', () => {
|
||||
expect(() => {
|
||||
stopQuotaMonitor();
|
||||
stopQuotaMonitor();
|
||||
stopQuotaMonitor();
|
||||
}).not.toThrow();
|
||||
});
|
||||
|
||||
it('should complete safely when called without prior start', () => {
|
||||
// No prior startQuotaMonitor call
|
||||
expect(() => {
|
||||
stopQuotaMonitor();
|
||||
}).not.toThrow();
|
||||
});
|
||||
|
||||
it('should handle multiple start/stop cycles', () => {
|
||||
const configDir = path.join(tmpDir, '.ccs', 'config');
|
||||
fs.mkdirSync(configDir, { recursive: true });
|
||||
fs.writeFileSync(
|
||||
path.join(configDir, 'unified-config.json'),
|
||||
JSON.stringify({
|
||||
version: 2,
|
||||
quota_management: {
|
||||
mode: 'auto',
|
||||
runtime_monitor: {
|
||||
enabled: false,
|
||||
normal_interval_seconds: 300,
|
||||
critical_interval_seconds: 60,
|
||||
warn_threshold: 20,
|
||||
exhaustion_threshold: 0,
|
||||
cooldown_minutes: 5,
|
||||
},
|
||||
},
|
||||
})
|
||||
);
|
||||
|
||||
expect(() => {
|
||||
startQuotaMonitor('agy', 'test@gmail.com');
|
||||
stopQuotaMonitor();
|
||||
startQuotaMonitor('agy', 'test@gmail.com');
|
||||
stopQuotaMonitor();
|
||||
}).not.toThrow();
|
||||
});
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user