mirror of
https://github.com/tiennm99/ccs.git
synced 2026-10-13 10:15:40 +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:
1 parent
0838b96429
commit
c6c94a0c1e
9 files changed
+788
No files matched your search
@@ -33,6 +33,7 @@ pnpm-lock.yaml
|
|||||||
package-lock.json
|
package-lock.json
|
||||||
|
|
||||||
.claude/active-plan
|
.claude/active-plan
|
||||||
|
.claude/agent-memory/
|
||||||
|
|
||||||
# Logs directory
|
# Logs directory
|
||||||
logs/
|
logs/
|
||||||
|
|||||||
@@ -374,3 +374,97 @@ export function maskEmail(email: string): string {
|
|||||||
function truncate(str: string, maxLen: number): string {
|
function truncate(str: string, maxLen: number): string {
|
||||||
return str.length > maxLen ? str.slice(0, maxLen - 3) + '...' : str;
|
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
|
// 13. Setup cleanup handlers
|
||||||
setupCleanupHandlers(
|
setupCleanupHandlers(
|
||||||
claude,
|
claude,
|
||||||
|
|||||||
@@ -21,6 +21,7 @@ import {
|
|||||||
import { detectRunningProxy, waitForProxyHealthy, reclaimOrphanedProxy } from '../proxy-detector';
|
import { detectRunningProxy, waitForProxyHealthy, reclaimOrphanedProxy } from '../proxy-detector';
|
||||||
import { withStartupLock } from '../startup-lock';
|
import { withStartupLock } from '../startup-lock';
|
||||||
import { killProcessOnPort } from '../../utils/platform-commands';
|
import { killProcessOnPort } from '../../utils/platform-commands';
|
||||||
|
import { stopQuotaMonitor } from '../quota-manager';
|
||||||
|
|
||||||
export interface ProxySessionResult {
|
export interface ProxySessionResult {
|
||||||
sessionId?: string;
|
sessionId?: string;
|
||||||
@@ -184,6 +185,7 @@ export function setupCleanupHandlers(
|
|||||||
};
|
};
|
||||||
|
|
||||||
const cleanup = () => {
|
const cleanup = () => {
|
||||||
|
stopQuotaMonitor();
|
||||||
log('Parent signal received, cleaning up');
|
log('Parent signal received, cleaning up');
|
||||||
|
|
||||||
if (
|
if (
|
||||||
@@ -214,6 +216,7 @@ export function setupCleanupHandlers(
|
|||||||
};
|
};
|
||||||
|
|
||||||
claude.on('exit', (code, signal) => {
|
claude.on('exit', (code, signal) => {
|
||||||
|
stopQuotaMonitor();
|
||||||
log(`Claude exited: code=${code}, signal=${signal}`);
|
log(`Claude exited: code=${code}, signal=${signal}`);
|
||||||
|
|
||||||
if (
|
if (
|
||||||
@@ -250,6 +253,7 @@ export function setupCleanupHandlers(
|
|||||||
});
|
});
|
||||||
|
|
||||||
claude.on('error', (error) => {
|
claude.on('error', (error) => {
|
||||||
|
stopQuotaMonitor();
|
||||||
console.error(require('../../utils/ui').fail(`Claude CLI error: ${error}`));
|
console.error(require('../../utils/ui').fail(`Claude CLI error: ${error}`));
|
||||||
|
|
||||||
if (
|
if (
|
||||||
|
|||||||
@@ -23,6 +23,7 @@ import {
|
|||||||
type AccountInfo,
|
type AccountInfo,
|
||||||
} from './account-manager';
|
} from './account-manager';
|
||||||
import { loadOrCreateUnifiedConfig } from '../config/unified-config-loader';
|
import { loadOrCreateUnifiedConfig } from '../config/unified-config-loader';
|
||||||
|
import type { RuntimeMonitorConfig } from '../config/unified-config-types';
|
||||||
|
|
||||||
// ============================================================================
|
// ============================================================================
|
||||||
// QUOTA CACHE (30-second TTL)
|
// QUOTA CACHE (30-second TTL)
|
||||||
@@ -416,3 +417,137 @@ export async function getQuotaStatus(provider: CLIProxyProvider): Promise<{
|
|||||||
|
|
||||||
return { accounts: results };
|
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 ??
|
partial.quota_management?.manual?.tier_lock ??
|
||||||
DEFAULT_QUOTA_MANAGEMENT_CONFIG.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 config - auto/manual/off control for reasoning budget
|
||||||
thinking: {
|
thinking: {
|
||||||
|
|||||||
@@ -369,6 +369,25 @@ export interface AutoQuotaConfig {
|
|||||||
cooldown_minutes: number;
|
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.
|
* Manual quota management configuration.
|
||||||
* User-controlled overrides for account selection.
|
* User-controlled overrides for account selection.
|
||||||
@@ -401,6 +420,8 @@ export interface QuotaManagementConfig {
|
|||||||
auto: AutoQuotaConfig;
|
auto: AutoQuotaConfig;
|
||||||
/** Manual mode settings */
|
/** Manual mode settings */
|
||||||
manual: ManualQuotaConfig;
|
manual: ManualQuotaConfig;
|
||||||
|
/** Runtime monitor settings */
|
||||||
|
runtime_monitor: RuntimeMonitorConfig;
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
@@ -422,6 +443,18 @@ export const DEFAULT_MANUAL_QUOTA_CONFIG: ManualQuotaConfig = {
|
|||||||
tier_lock: null,
|
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.
|
* Default quota management configuration.
|
||||||
*/
|
*/
|
||||||
@@ -429,6 +462,7 @@ export const DEFAULT_QUOTA_MANAGEMENT_CONFIG: QuotaManagementConfig = {
|
|||||||
mode: 'hybrid',
|
mode: 'hybrid',
|
||||||
auto: { ...DEFAULT_AUTO_QUOTA_CONFIG },
|
auto: { ...DEFAULT_AUTO_QUOTA_CONFIG },
|
||||||
manual: { ...DEFAULT_MANUAL_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