From c7058b42f4c2de9420563c46f4779dc3fa448e93 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 29 Jan 2026 23:05:10 +0000 Subject: [PATCH 1/7] feat(budget): add budget limits and cost controls for E2B sandboxes Implements comprehensive budget tracking and enforcement: - Add BudgetConfig, BudgetStatus, BudgetWarning types and BudgetExceededError class - Create ConfigManager for persistent settings (~/.parallel-cc/config.json) - Create BudgetTracker for period-based spending tracking (daily/weekly/monthly) - Add budget enforcement in SandboxManager with warnings at 50%/80% thresholds - Add CLI commands: config set/get/list, budget-status - Add --budget flag to sandbox-run command - Add budget checking during execution monitoring in ClaudeRunner - Add database methods for budget tracking records Test coverage: - 37 tests for ConfigManager - 38 tests for BudgetTracker - 25 tests for budget enforcement in SandboxManager https://claude.ai/code/session_01CDYWMZJdg8FH55bgQzNmds --- src/budget-tracker.ts | 345 ++++++++++++++++++++ src/cli.ts | 238 ++++++++++++++ src/config.ts | 248 ++++++++++++++ src/db.ts | 162 +++++++++- src/e2b/claude-runner.ts | 20 ++ src/e2b/sandbox-manager.ts | 154 ++++++++- src/types.ts | 65 ++++ tests/budget-tracker.test.ts | 462 +++++++++++++++++++++++++++ tests/config.test.ts | 426 ++++++++++++++++++++++++ tests/e2b/budget-enforcement.test.ts | 382 ++++++++++++++++++++++ 10 files changed, 2494 insertions(+), 8 deletions(-) create mode 100644 src/budget-tracker.ts create mode 100644 src/config.ts create mode 100644 tests/budget-tracker.test.ts create mode 100644 tests/config.test.ts create mode 100644 tests/e2b/budget-enforcement.test.ts diff --git a/src/budget-tracker.ts b/src/budget-tracker.ts new file mode 100644 index 0000000..655e70b --- /dev/null +++ b/src/budget-tracker.ts @@ -0,0 +1,345 @@ +/** + * Budget tracking and enforcement for parallel-cc + * + * Manages cost tracking by period, budget limits, and spending reports. + */ + +import type { SessionDB } from './db.js'; +import type { ConfigManager } from './config.js'; +import type { + BudgetPeriod, + BudgetTracking, + BudgetStatus +} from './types.js'; + +/** + * Result of a budget check + */ +export interface BudgetCheckResult { + allowed: boolean; + currentSpent: number; + estimatedCost: number; + remaining?: number; + limit?: number; + message?: string; +} + +/** + * Budget warning when threshold is reached + */ +export interface BudgetThresholdWarning { + threshold: number; + percentUsed: number; + currentSpent: number; + limit: number; + message: string; +} + +/** + * Result of session budget validation + */ +export interface SessionBudgetValidation { + valid: boolean; + message?: string; + warning?: string; +} + +/** + * BudgetTracker - Manages budget tracking and enforcement + * + * Features: + * - Period-based spending tracking (daily, weekly, monthly) + * - Budget limit enforcement with configurable thresholds + * - Session cost recording + * - Budget status reports + */ +export class BudgetTracker { + private db: SessionDB; + private configManager: ConfigManager; + private warningsIssued: Set = new Set(); + + constructor(db: SessionDB, configManager: ConfigManager) { + this.db = db; + this.configManager = configManager; + } + + /** + * Calculate the start date for a given period + * + * @param period - Budget period type + * @param date - Reference date (default: now) + * @returns ISO date string (YYYY-MM-DD) + */ + getPeriodStart(period: BudgetPeriod, date: Date = new Date()): string { + const year = date.getUTCFullYear(); + const month = date.getUTCMonth(); + const day = date.getUTCDate(); + + switch (period) { + case 'daily': + return `${year}-${String(month + 1).padStart(2, '0')}-${String(day).padStart(2, '0')}`; + + case 'weekly': { + // Get Monday of the current week + const dayOfWeek = date.getUTCDay(); + const daysFromMonday = dayOfWeek === 0 ? 6 : dayOfWeek - 1; + const monday = new Date(date); + monday.setUTCDate(day - daysFromMonday); + return `${monday.getUTCFullYear()}-${String(monday.getUTCMonth() + 1).padStart(2, '0')}-${String(monday.getUTCDate()).padStart(2, '0')}`; + } + + case 'monthly': + return `${year}-${String(month + 1).padStart(2, '0')}-01`; + + default: + throw new Error(`Unknown budget period: ${period}`); + } + } + + /** + * Get or create a budget tracking record for the current period + * + * @param period - Budget period type + * @returns Budget tracking record + */ + getOrCreatePeriodRecord(period: BudgetPeriod): BudgetTracking { + const periodStart = this.getPeriodStart(period); + const budgetConfig = this.configManager.getBudgetConfig(); + + // Get budget limit for this period type + let budgetLimit: number | undefined; + if (period === 'monthly' && budgetConfig.monthlyLimit !== undefined) { + budgetLimit = budgetConfig.monthlyLimit; + } + + return this.db.getOrCreateBudgetTrackingRecord(period, periodStart, budgetLimit); + } + + /** + * Record a cost to a period + * + * @param cost - Cost in USD + * @param period - Budget period type (default: monthly) + */ + recordCost(cost: number, period: BudgetPeriod = 'monthly'): void { + if (cost < 0) { + throw new Error('Cost must be a non-negative number'); + } + + if (cost === 0) { + return; + } + + const record = this.getOrCreatePeriodRecord(period); + this.db.updateBudgetSpent(record.id, cost); + } + + /** + * Record cost for a specific session + * + * Updates both the session's cost_estimate and the monthly budget tracking. + * + * @param sessionId - Session ID + * @param cost - Cost in USD + */ + recordSessionCost(sessionId: string, cost: number): void { + // Update session cost estimate + this.db.updateSessionCost(sessionId, cost); + + // Update monthly tracking + this.recordCost(cost, 'monthly'); + } + + /** + * Check if adding a cost would exceed the monthly budget + * + * @param estimatedCost - Estimated cost to add + * @returns Budget check result + */ + checkMonthlyBudget(estimatedCost: number): BudgetCheckResult { + const budgetConfig = this.configManager.getBudgetConfig(); + const monthlyLimit = budgetConfig.monthlyLimit; + + const currentSpent = this.getCurrentSpending('monthly'); + + // No limit configured + if (monthlyLimit === undefined || monthlyLimit === null) { + return { + allowed: true, + currentSpent, + estimatedCost + }; + } + + const remaining = monthlyLimit - currentSpent; + const wouldExceed = currentSpent + estimatedCost > monthlyLimit; + + return { + allowed: !wouldExceed, + currentSpent, + estimatedCost, + remaining, + limit: monthlyLimit, + message: wouldExceed + ? `Adding $${estimatedCost.toFixed(2)} would exceed monthly budget. ` + + `Current: $${currentSpent.toFixed(2)}, Limit: $${monthlyLimit.toFixed(2)}, ` + + `Remaining: $${remaining.toFixed(2)}` + : undefined + }; + } + + /** + * Get current spending for a period + * + * @param period - Budget period type + * @returns Current spending amount + */ + getCurrentSpending(period: BudgetPeriod): number { + const record = this.getOrCreatePeriodRecord(period); + return record.spent; + } + + /** + * Generate a budget status report + * + * @param period - Budget period type (default: monthly) + * @returns Budget status report + */ + generateBudgetStatus(period: BudgetPeriod = 'monthly'): BudgetStatus { + const periodRecord = this.getOrCreatePeriodRecord(period); + const budgetConfig = this.configManager.getBudgetConfig(); + + // Get E2B sessions with full data (including cost fields) + const e2bSessions = this.db.listE2BSessions(); + + // Build session info - get full session data to include cost fields + const sessions = e2bSessions.map(e2bSession => { + // Get full session data which includes cost fields + const fullSession = this.db.getSessionById(e2bSession.id); + return { + sessionId: e2bSession.id, + sandboxId: e2bSession.sandbox_id, + budgetLimit: fullSession?.budget_limit ?? undefined, + costEstimate: fullSession?.cost_estimate ?? undefined, + status: e2bSession.status ?? undefined, + createdAt: e2bSession.created_at + }; + }); + + // Calculate totals + const totalSpent = periodRecord.spent; + const limit = period === 'monthly' ? budgetConfig.monthlyLimit : undefined; + const remaining = limit !== undefined ? limit - totalSpent : undefined; + + return { + currentPeriod: { + period, + start: periodRecord.periodStart, + limit, + spent: totalSpent, + remaining + }, + sessions, + totalSpent, + remainingBudget: remaining + }; + } + + /** + * Get warning thresholds from config + * + * @returns Array of threshold percentages (e.g., [0.5, 0.8]) + */ + getWarningThresholds(): number[] { + const budgetConfig = this.configManager.getBudgetConfig(); + return budgetConfig.warningThresholds ?? [0.5, 0.8]; + } + + /** + * Check if any warning thresholds have been crossed + * + * @returns Warning info if threshold crossed, null otherwise + */ + checkWarningThresholds(): BudgetThresholdWarning | null { + const budgetConfig = this.configManager.getBudgetConfig(); + const monthlyLimit = budgetConfig.monthlyLimit; + + // No limit means no warnings + if (monthlyLimit === undefined || monthlyLimit === null || monthlyLimit === 0) { + return null; + } + + const currentSpent = this.getCurrentSpending('monthly'); + const percentUsed = currentSpent / monthlyLimit; + const thresholds = this.getWarningThresholds().sort((a, b) => b - a); // Sort descending + + // Find the highest threshold that has been crossed but not yet warned about + for (const threshold of thresholds) { + if (percentUsed >= threshold && !this.warningsIssued.has(threshold)) { + this.warningsIssued.add(threshold); + + return { + threshold, + percentUsed, + currentSpent, + limit: monthlyLimit, + message: `Budget warning: ${(percentUsed * 100).toFixed(0)}% of monthly limit used ` + + `($${currentSpent.toFixed(2)} / $${monthlyLimit.toFixed(2)})` + }; + } + } + + return null; + } + + /** + * Get the default per-session budget from config + * + * @returns Per-session default or undefined if not configured + */ + getPerSessionDefault(): number | undefined { + const budgetConfig = this.configManager.getBudgetConfig(); + return budgetConfig.perSessionDefault; + } + + /** + * Validate a session budget amount + * + * @param budget - Proposed session budget + * @returns Validation result + */ + validateSessionBudget(budget: number): SessionBudgetValidation { + if (budget < 0) { + return { + valid: false, + message: 'Session budget must be a positive number' + }; + } + + // Check if this would likely exceed the monthly budget + const budgetConfig = this.configManager.getBudgetConfig(); + const monthlyLimit = budgetConfig.monthlyLimit; + + if (monthlyLimit !== undefined && monthlyLimit !== null) { + const currentSpent = this.getCurrentSpending('monthly'); + const remaining = monthlyLimit - currentSpent; + + if (budget > remaining) { + return { + valid: true, // Still valid, but warn + warning: `Session budget ($${budget.toFixed(2)}) would exceed remaining monthly budget ` + + `($${remaining.toFixed(2)}). Total monthly limit: $${monthlyLimit.toFixed(2)}` + }; + } + } + + return { valid: true }; + } + + /** + * Reset warning tracking (for new sessions) + */ + resetWarnings(): void { + this.warningsIssued.clear(); + } +} diff --git a/src/cli.ts b/src/cli.ts index 6f978f9..cd14292 100644 --- a/src/cli.ts +++ b/src/cli.ts @@ -34,6 +34,8 @@ import { executeClaudeInSandbox } from './e2b/claude-runner.js'; import { pushToRemoteAndCreatePR } from './e2b/git-live.js'; import { validateSSHKeyPath, injectSSHKey, cleanupSSHKey, getSecurityWarning } from './e2b/ssh-key-injector.js'; import { TemplateManager, validateTemplateName, validateTemplate } from './e2b/templates.js'; +import { ConfigManager, DEFAULT_CONFIG_PATH } from './config.js'; +import { BudgetTracker } from './budget-tracker.js'; import { logger } from './logger.js'; import * as fs from 'fs/promises'; import { existsSync } from 'fs'; @@ -1437,6 +1439,7 @@ Examples: .option('--confirm-ssh-key', 'Skip interactive SSH key security warning (for non-interactive use)') .option('--npm-token ', 'NPM authentication token for private packages (or set PARALLEL_CC_NPM_TOKEN env var)') .option('--npm-registry ', 'Custom NPM registry URL (default: https://registry.npmjs.org)', 'https://registry.npmjs.org') + .option('--budget ', 'Per-session budget limit in USD (e.g., 0.50 for $0.50)') .option('--json', 'Output as JSON') .action(async (options) => { const coordinator = new Coordinator(); @@ -1749,6 +1752,18 @@ Examples: sandbox = createResult.sandbox; sandboxId = createResult.sandboxId; // Track for cleanup in catch block + // Set budget limit if provided (v1.1) + if (options.budget) { + const budgetAmount = parseFloat(options.budget); + if (isNaN(budgetAmount) || budgetAmount < 0) { + throw new Error(`Invalid budget amount: ${options.budget}. Must be a positive number.`); + } + sandboxManager.setBudgetLimit(sandboxId, budgetAmount); + if (!options.json) { + console.log(chalk.dim(` Budget limit: $${budgetAmount.toFixed(2)}`)); + } + } + if (!options.json) { console.log(chalk.green(`✓ Sandbox created: ${sandboxId}`)); console.log(chalk.dim(' Uploading workspace...')); @@ -3003,5 +3018,228 @@ templatesCmd } }); +// ============================================================================ +// Config Commands (v1.1) +// ============================================================================ + +/** + * Config command - Manage parallel-cc configuration + */ +const configCmd = program + .command('config') + .description('Manage parallel-cc configuration (v1.1)'); + +configCmd + .command('set ') + .description('Set a configuration value (e.g., budget.monthly-limit 10.00)') + .option('--json', 'Output as JSON') + .action(async (key: string, value: string, options) => { + try { + const configManager = new ConfigManager(); + + // Convert key from kebab-case to camelCase for budget keys + const normalizedKey = key.replace('monthly-limit', 'monthlyLimit') + .replace('per-session-default', 'perSessionDefault') + .replace('warning-thresholds', 'warningThresholds'); + + // Parse value based on key type + let parsedValue: unknown; + + if (normalizedKey.includes('budget')) { + // Budget values are numbers or arrays + if (normalizedKey.includes('Thresholds')) { + // Parse as array of numbers + parsedValue = value.split(',').map(v => parseFloat(v.trim())); + } else { + // Parse as number + parsedValue = parseFloat(value); + if (isNaN(parsedValue as number)) { + throw new Error(`Invalid number: ${value}`); + } + } + } else if (value === 'true' || value === 'false') { + parsedValue = value === 'true'; + } else if (!isNaN(parseFloat(value))) { + parsedValue = parseFloat(value); + } else { + parsedValue = value; + } + + configManager.set(normalizedKey, parsedValue); + + if (options.json) { + console.log(JSON.stringify({ success: true, key: normalizedKey, value: parsedValue })); + } else { + console.log(chalk.green(`✓ Set ${normalizedKey} = ${JSON.stringify(parsedValue)}`)); + } + } catch (error) { + const errorMessage = error instanceof Error ? error.message : 'Unknown error'; + if (options.json) { + console.log(JSON.stringify({ success: false, error: errorMessage })); + } else { + console.error(chalk.red(`✗ Failed to set config: ${errorMessage}`)); + } + process.exit(1); + } + }); + +configCmd + .command('get ') + .description('Get a configuration value') + .option('--json', 'Output as JSON') + .action(async (key: string, options) => { + try { + const configManager = new ConfigManager(); + + // Convert key from kebab-case to camelCase + const normalizedKey = key.replace('monthly-limit', 'monthlyLimit') + .replace('per-session-default', 'perSessionDefault') + .replace('warning-thresholds', 'warningThresholds'); + + const value = configManager.get(normalizedKey); + + if (options.json) { + console.log(JSON.stringify({ key: normalizedKey, value })); + } else { + if (value === undefined) { + console.log(chalk.yellow(`${normalizedKey}: (not set)`)); + } else { + console.log(`${normalizedKey}: ${JSON.stringify(value)}`); + } + } + } catch (error) { + const errorMessage = error instanceof Error ? error.message : 'Unknown error'; + if (options.json) { + console.log(JSON.stringify({ error: errorMessage })); + } else { + console.error(chalk.red(`✗ Failed to get config: ${errorMessage}`)); + } + process.exit(1); + } + }); + +configCmd + .command('list') + .description('Display all configuration values') + .option('--json', 'Output as JSON') + .action(async (options) => { + try { + const configManager = new ConfigManager(); + const config = configManager.getAll(); + + if (options.json) { + console.log(JSON.stringify(config, null, 2)); + } else { + console.log(chalk.bold('Configuration:')); + console.log(chalk.dim(` File: ${DEFAULT_CONFIG_PATH}\n`)); + + // Display budget config + console.log(chalk.cyan('Budget Settings:')); + const budget = config.budget; + console.log(` monthly-limit: ${budget.monthlyLimit !== undefined ? `$${budget.monthlyLimit.toFixed(2)}` : chalk.dim('(not set)')}`); + console.log(` per-session-default: ${budget.perSessionDefault !== undefined ? `$${budget.perSessionDefault.toFixed(2)}` : chalk.dim('(not set)')}`); + console.log(` warning-thresholds: ${budget.warningThresholds ? budget.warningThresholds.map(t => `${(t * 100).toFixed(0)}%`).join(', ') : chalk.dim('(not set)')}`); + + // Display other config keys + const otherKeys = Object.keys(config).filter(k => k !== 'budget'); + if (otherKeys.length > 0) { + console.log(chalk.cyan('\nOther Settings:')); + for (const key of otherKeys) { + console.log(` ${key}: ${JSON.stringify(config[key])}`); + } + } + } + } catch (error) { + const errorMessage = error instanceof Error ? error.message : 'Unknown error'; + if (options.json) { + console.log(JSON.stringify({ error: errorMessage })); + } else { + console.error(chalk.red(`✗ Failed to list config: ${errorMessage}`)); + } + process.exit(1); + } + }); + +// ============================================================================ +// Budget Status Command (v1.1) +// ============================================================================ + +/** + * Budget status command - Show spending and budget information + */ +program + .command('budget-status') + .description('Show budget and spending status (v1.1)') + .option('--period ', 'Show specific period (daily, weekly, monthly)', 'monthly') + .option('--json', 'Output as JSON') + .action(async (options) => { + try { + const db = new SessionDB(); + + // Run migration if needed + if (!db.hasBudgetTrackingTable()) { + console.log(chalk.yellow('Running database migration for budget tracking...')); + await db.migrateToLatest(); + } + + const configManager = new ConfigManager(); + const tracker = new BudgetTracker(db, configManager); + + const period = options.period as 'daily' | 'weekly' | 'monthly'; + const status = tracker.generateBudgetStatus(period); + + if (options.json) { + console.log(JSON.stringify(status, null, 2)); + } else { + console.log(chalk.bold(`\nBudget Status (${period})\n`)); + + // Current period info + console.log(chalk.cyan('Current Period:')); + console.log(` Start: ${status.currentPeriod.start}`); + console.log(` Spent: ${chalk.yellow(`$${status.currentPeriod.spent.toFixed(2)}`)}`); + + if (status.currentPeriod.limit !== undefined) { + console.log(` Limit: $${status.currentPeriod.limit.toFixed(2)}`); + const remaining = status.currentPeriod.remaining ?? 0; + const remainingColor = remaining > status.currentPeriod.limit * 0.2 ? chalk.green : chalk.yellow; + console.log(` Remaining: ${remainingColor(`$${remaining.toFixed(2)}`)}`); + const percentUsed = (status.currentPeriod.spent / status.currentPeriod.limit) * 100; + console.log(` Usage: ${percentUsed.toFixed(1)}%`); + } else { + console.log(chalk.dim(' Limit: (not set)')); + } + + // Active sessions + if (status.sessions.length > 0) { + console.log(chalk.cyan('\nActive E2B Sessions:')); + for (const session of status.sessions) { + const cost = session.costEstimate !== undefined ? `$${session.costEstimate.toFixed(2)}` : 'calculating...'; + const budget = session.budgetLimit !== undefined ? ` (limit: $${session.budgetLimit.toFixed(2)})` : ''; + console.log(` ${session.sessionId.substring(0, 8)}... - ${cost}${budget}`); + } + } else { + console.log(chalk.dim('\nNo active E2B sessions.')); + } + + // Summary + console.log(chalk.cyan('\nSummary:')); + console.log(` Total spent this ${period}: $${status.totalSpent.toFixed(2)}`); + if (status.remainingBudget !== undefined) { + console.log(` Remaining budget: $${status.remainingBudget.toFixed(2)}`); + } + } + + db.close(); + } catch (error) { + const errorMessage = error instanceof Error ? error.message : 'Unknown error'; + if (options.json) { + console.log(JSON.stringify({ error: errorMessage })); + } else { + console.error(chalk.red(`✗ Failed to get budget status: ${errorMessage}`)); + } + process.exit(1); + } + }); + // Parse and execute program.parse(); diff --git a/src/config.ts b/src/config.ts new file mode 100644 index 0000000..9a3ab74 --- /dev/null +++ b/src/config.ts @@ -0,0 +1,248 @@ +/** + * Configuration management for parallel-cc + * + * Handles persistent user settings stored in ~/.parallel-cc/config.json + */ + +import { existsSync, mkdirSync, readFileSync, writeFileSync } from 'fs'; +import { dirname, join as pathJoin } from 'path'; +import { homedir } from 'os'; +import type { BudgetConfig } from './types.js'; + +/** + * Default budget configuration values + */ +export const DEFAULT_BUDGET_CONFIG: BudgetConfig = { + monthlyLimit: undefined, + perSessionDefault: undefined, + warningThresholds: [0.5, 0.8] +}; + +/** + * Default configuration structure + */ +interface Config { + budget: BudgetConfig; + [key: string]: unknown; +} + +const DEFAULT_CONFIG: Config = { + budget: DEFAULT_BUDGET_CONFIG +}; + +/** + * Default config file path + */ +export const DEFAULT_CONFIG_PATH = pathJoin(homedir(), '.parallel-cc', 'config.json'); + +/** + * ConfigManager - Manages persistent user configuration + * + * Features: + * - JSON file storage with automatic directory creation + * - Dot notation support for nested keys (e.g., "budget.monthlyLimit") + * - Validation for budget-related settings + * - Automatic persistence on changes + */ +export class ConfigManager { + private configPath: string; + private config: Config; + + /** + * Create a new ConfigManager + * + * @param configPath - Path to config file (default: ~/.parallel-cc/config.json) + */ + constructor(configPath: string = DEFAULT_CONFIG_PATH) { + this.configPath = this.resolvePath(configPath); + this.ensureDirectory(); + this.config = this.load(); + } + + /** + * Resolve tilde in path to home directory + */ + private resolvePath(path: string): string { + if (path.startsWith('~')) { + return path.replace('~', homedir()); + } + return path; + } + + /** + * Ensure config directory exists + */ + private ensureDirectory(): void { + const dir = dirname(this.configPath); + if (!existsSync(dir)) { + mkdirSync(dir, { recursive: true }); + } + } + + /** + * Load config from file + */ + private load(): Config { + if (!existsSync(this.configPath)) { + // Create default config + this.save(structuredClone(DEFAULT_CONFIG)); + return structuredClone(DEFAULT_CONFIG); + } + + try { + const content = readFileSync(this.configPath, 'utf-8'); + const parsed = JSON.parse(content); + + // Merge with defaults to ensure all required fields exist + return { + ...structuredClone(DEFAULT_CONFIG), + ...parsed, + budget: { + ...DEFAULT_BUDGET_CONFIG, + ...parsed.budget + } + }; + } catch { + // Invalid JSON - return defaults + return structuredClone(DEFAULT_CONFIG); + } + } + + /** + * Save config to file + */ + private save(config?: Config): void { + const toSave = config ?? this.config; + writeFileSync(this.configPath, JSON.stringify(toSave, null, 2)); + } + + /** + * Get a config value by key (supports dot notation) + * + * @param key - Config key (e.g., "budget.monthlyLimit") + * @returns Config value or undefined if not found + */ + get(key: string): unknown { + const parts = key.split('.'); + let current: unknown = this.config; + + for (const part of parts) { + if (current === null || current === undefined) { + return undefined; + } + if (typeof current !== 'object') { + return undefined; + } + current = (current as Record)[part]; + } + + return current; + } + + /** + * Set a config value by key (supports dot notation) + * + * @param key - Config key (e.g., "budget.monthlyLimit") + * @param value - Value to set + */ + set(key: string, value: unknown): void { + const parts = key.split('.'); + let current: Record = this.config; + + // Navigate/create path to parent + for (let i = 0; i < parts.length - 1; i++) { + const part = parts[i]; + if (current[part] === undefined || current[part] === null || typeof current[part] !== 'object') { + current[part] = {}; + } + current = current[part] as Record; + } + + // Set the value + const lastPart = parts[parts.length - 1]; + current[lastPart] = value; + + this.save(); + } + + /** + * Delete a config key (supports dot notation) + * + * @param key - Config key to delete + */ + delete(key: string): void { + const parts = key.split('.'); + let current: Record = this.config; + + // Navigate to parent + for (let i = 0; i < parts.length - 1; i++) { + const part = parts[i]; + if (current[part] === undefined || typeof current[part] !== 'object') { + return; // Path doesn't exist, nothing to delete + } + current = current[part] as Record; + } + + // Delete the key + const lastPart = parts[parts.length - 1]; + delete current[lastPart]; + + this.save(); + } + + /** + * Get entire config object (returns a copy) + * + * @returns Copy of entire config + */ + getAll(): Config { + return structuredClone(this.config); + } + + /** + * Get budget configuration + * + * @returns Budget config + */ + getBudgetConfig(): BudgetConfig { + return structuredClone(this.config.budget); + } + + /** + * Set budget configuration (partial update) + * + * @param config - Partial budget config to merge + * @throws Error if validation fails + */ + setBudgetConfig(config: Partial): void { + // Validate budget limits + if (config.monthlyLimit !== undefined && config.monthlyLimit !== null) { + if (config.monthlyLimit < 0) { + throw new Error('Budget limit must be a positive number'); + } + } + + if (config.perSessionDefault !== undefined && config.perSessionDefault !== null) { + if (config.perSessionDefault < 0) { + throw new Error('Per-session budget must be a positive number'); + } + } + + // Validate warning thresholds + if (config.warningThresholds !== undefined) { + for (const threshold of config.warningThresholds) { + if (threshold < 0 || threshold > 1) { + throw new Error('Warning thresholds must be between 0 and 1'); + } + } + } + + // Merge config + this.config.budget = { + ...this.config.budget, + ...config + }; + + this.save(); + } +} diff --git a/src/db.ts b/src/db.ts index 2551807..3e4ce5c 100644 --- a/src/db.ts +++ b/src/db.ts @@ -30,7 +30,10 @@ import type { SuggestionFilters, E2BSession, E2BSessionRow, - ExecutionMode + ExecutionMode, + BudgetTracking, + BudgetTrackingRow, + BudgetPeriod } from './types.js'; import { SandboxStatus } from './types.js'; import { DEFAULT_CONFIG } from './types.js'; @@ -622,6 +625,21 @@ export class SessionDB { } } + /** + * Check if budget_tracking table exists (v1.1) + */ + hasBudgetTrackingTable(): boolean { + try { + const result = this.db.prepare(` + SELECT name FROM sqlite_master + WHERE type='table' AND name='budget_tracking' + `).get() as { name: string } | undefined; + return result !== undefined; + } catch { + return false; + } + } + // ============================================================================ // File Claims (v0.5) // ============================================================================ @@ -1599,6 +1617,148 @@ export class SessionDB { } } + // ============================================================================ + // Budget Tracking (v1.1) + // ============================================================================ + + /** + * Create a new budget tracking record + * + * @param period - Budget period type (daily, weekly, monthly) + * @param periodStart - ISO date string for period start + * @param budgetLimit - Optional budget limit for this period + * @returns The created budget tracking record + */ + createBudgetTrackingRecord( + period: BudgetPeriod, + periodStart: string, + budgetLimit?: number + ): BudgetTracking { + const id = randomUUID(); + const stmt = this.db.prepare(` + INSERT INTO budget_tracking (id, period, period_start, budget_limit, spent) + VALUES (?, ?, ?, ?, 0) + RETURNING * + `); + + const row = stmt.get(id, period, periodStart, budgetLimit ?? null) as BudgetTrackingRow; + return this.rowToBudgetTracking(row); + } + + /** + * Get budget tracking record by period and start date + * + * @param period - Budget period type + * @param periodStart - ISO date string for period start + * @returns Budget tracking record or null if not found + */ + getBudgetTrackingRecord(period: BudgetPeriod, periodStart: string): BudgetTracking | null { + const stmt = this.db.prepare(` + SELECT * FROM budget_tracking + WHERE period = ? AND period_start = ? + `); + const row = stmt.get(period, periodStart) as BudgetTrackingRow | undefined; + return row ? this.rowToBudgetTracking(row) : null; + } + + /** + * Update spent amount in budget tracking record + * + * @param id - Budget tracking record ID + * @param amount - Amount to add to spent (positive number) + * @returns true if record was updated + */ + updateBudgetSpent(id: string, amount: number): boolean { + const stmt = this.db.prepare(` + UPDATE budget_tracking + SET spent = spent + ? + WHERE id = ? + `); + const result = stmt.run(amount, id); + return result.changes > 0; + } + + /** + * Get or create budget tracking record for period + * + * @param period - Budget period type + * @param periodStart - ISO date string for period start + * @param budgetLimit - Optional budget limit for new records + * @returns Budget tracking record + */ + getOrCreateBudgetTrackingRecord( + period: BudgetPeriod, + periodStart: string, + budgetLimit?: number + ): BudgetTracking { + const existing = this.getBudgetTrackingRecord(period, periodStart); + if (existing) { + return existing; + } + return this.createBudgetTrackingRecord(period, periodStart, budgetLimit); + } + + /** + * Update session cost estimate + * + * @param sessionId - Session ID + * @param costEstimate - Estimated cost in USD + * @param actualCost - Optional actual cost in USD + * @returns true if session was updated + */ + updateSessionCost( + sessionId: string, + costEstimate: number, + actualCost?: number + ): boolean { + let query = `UPDATE sessions SET cost_estimate = ?`; + const params: unknown[] = [costEstimate]; + + if (actualCost !== undefined) { + query += `, actual_cost = ?`; + params.push(actualCost); + } + + query += ` WHERE id = ?`; + params.push(sessionId); + + const stmt = this.db.prepare(query); + const result = stmt.run(...params); + return result.changes > 0; + } + + /** + * Get all budget tracking records for a period type + * + * @param period - Budget period type + * @param limit - Maximum number of records to return (default: 10) + * @returns Array of budget tracking records, ordered by period_start DESC + */ + getBudgetTrackingHistory(period: BudgetPeriod, limit: number = 10): BudgetTracking[] { + const stmt = this.db.prepare(` + SELECT * FROM budget_tracking + WHERE period = ? + ORDER BY period_start DESC + LIMIT ? + `); + const rows = stmt.all(period, limit) as BudgetTrackingRow[]; + return rows.map(row => this.rowToBudgetTracking(row)); + } + + /** + * Convert budget tracking row to model + */ + private rowToBudgetTracking(row: BudgetTrackingRow): BudgetTracking { + return { + id: row.id, + period: row.period, + periodStart: row.period_start, + budgetLimit: row.budget_limit ?? undefined, + spent: row.spent, + createdAt: row.created_at + }; + } + close(): void { this.db.close(); } diff --git a/src/e2b/claude-runner.ts b/src/e2b/claude-runner.ts index 7062762..7bc0824 100644 --- a/src/e2b/claude-runner.ts +++ b/src/e2b/claude-runner.ts @@ -1439,6 +1439,26 @@ export async function monitorExecution( } } + // Check budget limit (v1.1) + try { + const budgetWarning = await sandboxManager.checkBudgetLimit(sandboxId); + if (budgetWarning) { + // Soft budget warning - log but continue + logger.warn(`Budget warning: ${budgetWarning.message}`); + } + } catch (budgetError) { + // BudgetExceededError is thrown when budget is exceeded + if (budgetError instanceof Error && budgetError.name === 'BudgetExceededError') { + logger.error(`Budget exceeded for sandbox ${sandboxId}`); + return { + shouldTerminate: true, + reason: budgetError.message + }; + } + // Other errors - log but don't terminate + logger.warn(`Budget check error: ${budgetError instanceof Error ? budgetError.message : String(budgetError)}`); + } + // Check sandbox health const healthCheck = await sandboxManager.monitorSandboxHealth(sandboxId); if (!healthCheck.isHealthy) { diff --git a/src/e2b/sandbox-manager.ts b/src/e2b/sandbox-manager.ts index 1c48614..1bf6f16 100644 --- a/src/e2b/sandbox-manager.ts +++ b/src/e2b/sandbox-manager.ts @@ -11,11 +11,13 @@ import { Sandbox } from 'e2b'; import type { Logger } from '../logger.js'; import { SandboxStatus, + BudgetExceededError, type E2BSessionConfig, type SandboxHealthCheck, type SandboxTerminationResult, type TimeoutWarning, - type SandboxTemplate + type SandboxTemplate, + type BudgetWarning } from '../types.js'; /** @@ -29,14 +31,20 @@ export interface TemplateApplicationResult { error?: string; } +// Extended config type with budget thresholds +interface ExtendedE2BSessionConfig extends E2BSessionConfig { + budgetWarningThresholds?: number[]; +} + // Default configuration const envTemplate = process.env.E2B_TEMPLATE?.trim(); -const DEFAULT_CONFIG: Required = { +const DEFAULT_CONFIG: Required = { claudeVersion: 'latest', e2bSdkVersion: '1.13.2', sandboxImage: (envTemplate && envTemplate.length > 0) ? envTemplate : 'anthropic-claude-code', // E2B template with pre-installed Claude Code timeoutMinutes: 60, - warningThresholds: [30, 50] + warningThresholds: [30, 50], + budgetWarningThresholds: [0.5, 0.8] // Default: warn at 50% and 80% of budget }; // Security constants @@ -104,13 +112,16 @@ export function validateFilePath(path: string): boolean { * - Graceful termination and cleanup */ export class SandboxManager { - private config: Required; + private config: Required; private logger: Logger; private activeSandboxes: Map = new Map(); private sandboxStartTimes: Map = new Map(); private timeoutWarningsIssued: Map> = new Map(); + // Budget tracking (v1.1) + private budgetLimits: Map = new Map(); + private budgetWarningsIssued: Map> = new Map(); - constructor(logger: Logger, config: Partial = {}) { + constructor(logger: Logger, config: Partial = {}) { this.logger = logger; this.config = { ...DEFAULT_CONFIG, ...config }; } @@ -335,6 +346,131 @@ export class SandboxManager { } } + // ============================================================================ + // Budget Enforcement (v1.1) + // ============================================================================ + + /** + * Set budget limit for a sandbox + * + * @param sandboxId - Sandbox ID + * @param limit - Budget limit in USD + * @throws Error if limit is negative + */ + setBudgetLimit(sandboxId: string, limit: number): void { + if (limit < 0) { + throw new Error('Budget limit must be a positive number'); + } + this.budgetLimits.set(sandboxId, limit); + this.budgetWarningsIssued.set(sandboxId, new Set()); + this.logger.info(`Set budget limit $${limit.toFixed(2)} for sandbox ${sandboxId}`); + } + + /** + * Get budget limit for a sandbox + * + * @param sandboxId - Sandbox ID + * @returns Budget limit or undefined if not set + */ + getBudgetLimit(sandboxId: string): number | undefined { + return this.budgetLimits.get(sandboxId); + } + + /** + * Check budget limit and issue warnings or terminate if exceeded + * + * @param sandboxId - Sandbox ID to check + * @returns Budget warning if threshold reached, null otherwise + * @throws BudgetExceededError if budget is exceeded + */ + async checkBudgetLimit(sandboxId: string): Promise { + try { + const budgetLimit = this.budgetLimits.get(sandboxId); + + // No budget limit set or zero (disabled) + if (budgetLimit === undefined || budgetLimit === 0) { + return null; + } + + const startTime = this.sandboxStartTimes.get(sandboxId); + if (!startTime) { + this.logger.warn(`Start time not found for sandbox ${sandboxId}, cannot check budget`); + return null; + } + + // Calculate current cost + const elapsedMinutes = Math.floor((Date.now() - startTime.getTime()) / 60000); + const currentCost = this.calculateEstimatedCostNumeric(elapsedMinutes); + const percentUsed = currentCost / budgetLimit; + + const warningsIssued = this.budgetWarningsIssued.get(sandboxId) || new Set(); + + // Hard limit enforcement (budget exceeded) + if (currentCost >= budgetLimit) { + const warning: BudgetWarning = { + sandboxId, + currentCost, + budgetLimit, + percentUsed, + warningLevel: 'hard', + message: `BUDGET EXCEEDED: Sandbox has used $${currentCost.toFixed(2)} of $${budgetLimit.toFixed(2)} budget (${(percentUsed * 100).toFixed(0)}%). Terminating sandbox.` + }; + + this.logger.error(warning.message); + + // Terminate sandbox + await this.terminateSandbox(sandboxId); + + throw new BudgetExceededError(sandboxId, currentCost, budgetLimit); + } + + // Check soft warning thresholds + const thresholds = this.config.budgetWarningThresholds.sort((a, b) => b - a); // Sort descending + + for (const threshold of thresholds) { + if (percentUsed >= threshold && !warningsIssued.has(threshold)) { + warningsIssued.add(threshold); + this.budgetWarningsIssued.set(sandboxId, warningsIssued); + + const warning: BudgetWarning = { + sandboxId, + currentCost, + budgetLimit, + percentUsed, + warningLevel: 'soft', + message: `Budget warning: ${(percentUsed * 100).toFixed(0)}% of budget used ($${currentCost.toFixed(2)} / $${budgetLimit.toFixed(2)})` + }; + + this.logger.warn(warning.message); + return warning; + } + } + + return null; + } catch (error) { + // Re-throw BudgetExceededError + if (error instanceof BudgetExceededError) { + throw error; + } + + const errorMsg = error instanceof Error ? error.message : String(error); + this.logger.error(`Budget check failed for sandbox ${sandboxId}: ${errorMsg}`); + return null; + } + } + + /** + * Calculate estimated cost as a number (for budget calculations) + * + * @param elapsedMinutes - Elapsed minutes + * @returns Estimated cost in USD + */ + private calculateEstimatedCostNumeric(elapsedMinutes: number): number { + // E2B pricing: ~$0.10/hour for basic compute + const costPerMinute = 0.10 / 60; + return elapsedMinutes * costPerMinute; + } + /** * Terminate a sandbox and cleanup resources * @@ -376,10 +512,12 @@ export class SandboxManager { } } - // Cleanup tracking data + // Cleanup tracking data (including budget) this.activeSandboxes.delete(sandboxId); this.sandboxStartTimes.delete(sandboxId); this.timeoutWarningsIssued.delete(sandboxId); + this.budgetLimits.delete(sandboxId); + this.budgetWarningsIssued.delete(sandboxId); this.logger.info(`E2B sandbox terminated successfully: ${sandboxId}`); @@ -392,10 +530,12 @@ export class SandboxManager { const errorMsg = error instanceof Error ? error.message : String(error); this.logger.error(`Failed to terminate sandbox ${sandboxId}: ${errorMsg}`); - // Best-effort cleanup even on error + // Best-effort cleanup even on error (including budget) this.activeSandboxes.delete(sandboxId); this.sandboxStartTimes.delete(sandboxId); this.timeoutWarningsIssued.delete(sandboxId); + this.budgetLimits.delete(sandboxId); + this.budgetWarningsIssued.delete(sandboxId); return { success: false, diff --git a/src/types.ts b/src/types.ts index 0c2268b..1479fb1 100644 --- a/src/types.ts +++ b/src/types.ts @@ -497,6 +497,7 @@ export interface E2BSessionConfig { sandboxImage: string; // Recommended: "anthropic-claude-code" (pre-installed Claude Code). Fallback: "base" or custom images auto-install via apt-get (Debian/Ubuntu required) then npm install -g @anthropic-ai/claude-code timeoutMinutes?: number; // Default: 60 warningThresholds?: number[]; // Default: [30, 50] minutes + budgetWarningThresholds?: number[]; // Default: [0.5, 0.8] (50%, 80% of budget) } /** @@ -611,6 +612,70 @@ export interface BudgetTracking { createdAt: string; } +/** + * Budget configuration for user settings + */ +export interface BudgetConfig { + /** Monthly spending limit in USD */ + monthlyLimit?: number; + /** Default per-session budget in USD */ + perSessionDefault?: number; + /** Warning thresholds as percentages (e.g., [0.5, 0.8] for 50% and 80%) */ + warningThresholds?: number[]; +} + +/** + * Budget status report for CLI command + */ +export interface BudgetStatus { + currentPeriod: { + period: BudgetPeriod; + start: string; + limit?: number; + spent: number; + remaining?: number; + }; + sessions: Array<{ + sessionId: string; + sandboxId?: string; + budgetLimit?: number; + costEstimate?: number; + status?: string; + createdAt: string; + }>; + totalSpent: number; + remainingBudget?: number; +} + +/** + * Budget warning during sandbox execution (similar to TimeoutWarning) + */ +export interface BudgetWarning { + sandboxId: string; + currentCost: number; + budgetLimit: number; + percentUsed: number; + warningLevel: 'soft' | 'hard'; + message: string; +} + +/** + * Error thrown when budget is exceeded + */ +export class BudgetExceededError extends Error { + public readonly sandboxId: string; + public readonly currentCost: number; + public readonly budgetLimit: number; + + constructor(sandboxId: string, currentCost: number, budgetLimit: number) { + super(`Budget exceeded for sandbox ${sandboxId}: $${currentCost.toFixed(2)} >= $${budgetLimit.toFixed(2)}`); + this.name = 'BudgetExceededError'; + this.sandboxId = sandboxId; + this.currentCost = currentCost; + this.budgetLimit = budgetLimit; + } +} + // ============================================================================ // Sandbox Template Types (v1.1) // ============================================================================ diff --git a/tests/budget-tracker.test.ts b/tests/budget-tracker.test.ts new file mode 100644 index 0000000..1688451 --- /dev/null +++ b/tests/budget-tracker.test.ts @@ -0,0 +1,462 @@ +/** + * Tests for BudgetTracker class + * + * TDD: These tests define the expected behavior of the BudgetTracker + * before implementation. + */ + +import { describe, it, expect, beforeEach, afterEach } from 'vitest'; +import * as fs from 'fs'; +import * as path from 'path'; +import * as os from 'os'; +import { SessionDB } from '../src/db.js'; +import { BudgetTracker } from '../src/budget-tracker.js'; +import { ConfigManager } from '../src/config.js'; +import type { BudgetPeriod } from '../src/types.js'; + +// Test fixtures directory - unique per process to avoid conflicts +const TEST_DIR = path.join(os.tmpdir(), 'parallel-cc-budget-test-' + process.pid); +const TEST_DB_PATH = path.join(TEST_DIR, 'test.db'); +const TEST_CONFIG_PATH = path.join(TEST_DIR, 'config.json'); + +describe('BudgetTracker', () => { + let db: SessionDB; + let configManager: ConfigManager; + let tracker: BudgetTracker; + + beforeEach(async () => { + // Create test directory + fs.mkdirSync(TEST_DIR, { recursive: true }); + db = new SessionDB(TEST_DB_PATH); + + // Run migrations to get budget_tracking table + await db.migrateToLatest(); + + configManager = new ConfigManager(TEST_CONFIG_PATH); + tracker = new BudgetTracker(db, configManager); + }); + + afterEach(() => { + // Close database and clean up + db.close(); + fs.rmSync(TEST_DIR, { recursive: true, force: true }); + }); + + // ========================================================================== + // Period Calculation Tests + // ========================================================================== + + describe('getPeriodStart', () => { + it('should calculate correct daily period start', () => { + const now = new Date('2025-01-15T14:30:00Z'); + const start = tracker.getPeriodStart('daily', now); + + expect(start).toBe('2025-01-15'); + }); + + it('should calculate correct weekly period start (Monday)', () => { + const wednesday = new Date('2025-01-15T14:30:00Z'); // Wednesday + const start = tracker.getPeriodStart('weekly', wednesday); + + // Monday of that week + expect(start).toBe('2025-01-13'); + }); + + it('should calculate correct monthly period start', () => { + const midMonth = new Date('2025-01-15T14:30:00Z'); + const start = tracker.getPeriodStart('monthly', midMonth); + + expect(start).toBe('2025-01-01'); + }); + + it('should handle month boundaries correctly', () => { + const lastDayOfMonth = new Date('2025-01-31T23:59:59Z'); + const start = tracker.getPeriodStart('monthly', lastDayOfMonth); + + expect(start).toBe('2025-01-01'); + }); + }); + + // ========================================================================== + // Budget Tracking Record Tests + // ========================================================================== + + describe('getOrCreatePeriodRecord', () => { + it('should create new period record if none exists', () => { + const record = tracker.getOrCreatePeriodRecord('monthly'); + + expect(record).toBeDefined(); + expect(record.period).toBe('monthly'); + expect(record.spent).toBe(0); + }); + + it('should return existing period record', () => { + const record1 = tracker.getOrCreatePeriodRecord('monthly'); + record1.spent = 5.00; + + // Update spent in the record + tracker.recordCost(5.00, 'monthly'); + + const record2 = tracker.getOrCreatePeriodRecord('monthly'); + + expect(record2.id).toBe(record1.id); + expect(record2.spent).toBe(5.00); + }); + + it('should create separate records for different periods', () => { + const daily = tracker.getOrCreatePeriodRecord('daily'); + const monthly = tracker.getOrCreatePeriodRecord('monthly'); + + expect(daily.id).not.toBe(monthly.id); + expect(daily.period).toBe('daily'); + expect(monthly.period).toBe('monthly'); + }); + }); + + // ========================================================================== + // Recording Costs Tests + // ========================================================================== + + describe('recordCost', () => { + it('should record cost to current period', () => { + tracker.recordCost(0.50, 'monthly'); + + const record = tracker.getOrCreatePeriodRecord('monthly'); + expect(record.spent).toBe(0.50); + }); + + it('should accumulate costs', () => { + tracker.recordCost(0.25, 'monthly'); + tracker.recordCost(0.35, 'monthly'); + tracker.recordCost(0.40, 'monthly'); + + const record = tracker.getOrCreatePeriodRecord('monthly'); + expect(record.spent).toBeCloseTo(1.00, 2); + }); + + it('should track costs separately per period type', () => { + tracker.recordCost(1.00, 'daily'); + tracker.recordCost(5.00, 'monthly'); + + const daily = tracker.getOrCreatePeriodRecord('daily'); + const monthly = tracker.getOrCreatePeriodRecord('monthly'); + + expect(daily.spent).toBe(1.00); + expect(monthly.spent).toBe(5.00); + }); + + it('should throw on negative cost', () => { + expect(() => tracker.recordCost(-0.50, 'monthly')) + .toThrow('Cost must be a non-negative number'); + }); + + it('should allow zero cost (no-op)', () => { + expect(() => tracker.recordCost(0, 'monthly')).not.toThrow(); + + const record = tracker.getOrCreatePeriodRecord('monthly'); + expect(record.spent).toBe(0); + }); + }); + + // ========================================================================== + // recordSessionCost Tests + // ========================================================================== + + describe('recordSessionCost', () => { + it('should update session cost estimate', async () => { + // Create a test E2B session first + const session = db.createE2BSession({ + id: 'test-session-1', + pid: 12345, + repo_path: '/test/repo', + worktree_path: '/test/worktree', + worktree_name: 'test-branch', + sandbox_id: 'sandbox-123', + prompt: 'Test prompt' + }); + + tracker.recordSessionCost(session.id, 0.75); + + const updated = db.getSessionById(session.id); + expect(updated?.cost_estimate).toBe(0.75); + }); + + it('should update both session and period tracking', async () => { + const session = db.createE2BSession({ + id: 'test-session-2', + pid: 12346, + repo_path: '/test/repo', + worktree_path: '/test/worktree', + worktree_name: 'test-branch', + sandbox_id: 'sandbox-456', + prompt: 'Test prompt' + }); + + tracker.recordSessionCost(session.id, 1.25); + + const updated = db.getSessionById(session.id); + expect(updated?.cost_estimate).toBe(1.25); + + const monthlyRecord = tracker.getOrCreatePeriodRecord('monthly'); + expect(monthlyRecord.spent).toBe(1.25); + }); + }); + + // ========================================================================== + // Budget Check Tests + // ========================================================================== + + describe('checkMonthlyBudget', () => { + it('should return true when no monthly limit set', () => { + const result = tracker.checkMonthlyBudget(10.00); + expect(result.allowed).toBe(true); + }); + + it('should return true when under budget', () => { + configManager.setBudgetConfig({ monthlyLimit: 10.00 }); + + const result = tracker.checkMonthlyBudget(0.50); + expect(result.allowed).toBe(true); + expect(result.remaining).toBe(10.00); + }); + + it('should return false when would exceed budget', () => { + configManager.setBudgetConfig({ monthlyLimit: 1.00 }); + + // Already spent 0.75 + tracker.recordCost(0.75, 'monthly'); + + // Trying to spend 0.50 more would exceed + const result = tracker.checkMonthlyBudget(0.50); + expect(result.allowed).toBe(false); + expect(result.remaining).toBe(0.25); + expect(result.message).toContain('exceed'); + }); + + it('should return true when exactly at limit', () => { + configManager.setBudgetConfig({ monthlyLimit: 1.00 }); + tracker.recordCost(0.75, 'monthly'); + + // Spending exactly to the limit + const result = tracker.checkMonthlyBudget(0.25); + expect(result.allowed).toBe(true); + }); + + it('should include spending details in result', () => { + configManager.setBudgetConfig({ monthlyLimit: 10.00 }); + tracker.recordCost(3.00, 'monthly'); + + const result = tracker.checkMonthlyBudget(2.00); + + expect(result.currentSpent).toBe(3.00); + expect(result.estimatedCost).toBe(2.00); + expect(result.remaining).toBe(7.00); + expect(result.limit).toBe(10.00); + }); + }); + + // ========================================================================== + // Budget Status Report Tests + // ========================================================================== + + describe('generateBudgetStatus', () => { + it('should generate status report with no sessions', () => { + const status = tracker.generateBudgetStatus('monthly'); + + expect(status.currentPeriod.period).toBe('monthly'); + expect(status.currentPeriod.spent).toBe(0); + expect(status.sessions).toEqual([]); + expect(status.totalSpent).toBe(0); + }); + + it('should include active E2B sessions', async () => { + db.createE2BSession({ + id: 'session-1', + pid: 12345, + repo_path: '/test/repo', + worktree_path: '/test/worktree', + worktree_name: null, + sandbox_id: 'sandbox-1', + prompt: 'Task 1' + }); + + db.createE2BSession({ + id: 'session-2', + pid: 12346, + repo_path: '/test/repo', + worktree_path: '/test/worktree2', + worktree_name: null, + sandbox_id: 'sandbox-2', + prompt: 'Task 2' + }); + + const status = tracker.generateBudgetStatus('monthly'); + + expect(status.sessions).toHaveLength(2); + expect(status.sessions[0].sessionId).toBeDefined(); + expect(status.sessions[0].sandboxId).toBeDefined(); + }); + + it('should include budget limit from config', () => { + configManager.setBudgetConfig({ monthlyLimit: 25.00 }); + + const status = tracker.generateBudgetStatus('monthly'); + + expect(status.currentPeriod.limit).toBe(25.00); + expect(status.remainingBudget).toBe(25.00); + }); + + it('should calculate remaining budget correctly', () => { + configManager.setBudgetConfig({ monthlyLimit: 10.00 }); + tracker.recordCost(3.50, 'monthly'); + + const status = tracker.generateBudgetStatus('monthly'); + + expect(status.currentPeriod.spent).toBe(3.50); + expect(status.currentPeriod.remaining).toBe(6.50); + expect(status.totalSpent).toBe(3.50); + expect(status.remainingBudget).toBe(6.50); + }); + + it('should include session cost estimates', async () => { + const session = db.createE2BSession({ + id: 'session-cost', + pid: 12347, + repo_path: '/test/repo', + worktree_path: '/test/worktree', + worktree_name: null, + sandbox_id: 'sandbox-cost', + prompt: 'Cost test' + }); + + tracker.recordSessionCost(session.id, 0.75); + + const status = tracker.generateBudgetStatus('monthly'); + + const sessionInfo = status.sessions.find(s => s.sessionId === session.id); + expect(sessionInfo?.costEstimate).toBe(0.75); + }); + }); + + // ========================================================================== + // getCurrentSpending Tests + // ========================================================================== + + describe('getCurrentSpending', () => { + it('should return 0 for fresh tracker', () => { + expect(tracker.getCurrentSpending('monthly')).toBe(0); + expect(tracker.getCurrentSpending('daily')).toBe(0); + }); + + it('should return accumulated spending', () => { + tracker.recordCost(1.00, 'monthly'); + tracker.recordCost(2.00, 'monthly'); + + expect(tracker.getCurrentSpending('monthly')).toBe(3.00); + }); + + it('should track different periods independently', () => { + tracker.recordCost(1.00, 'daily'); + tracker.recordCost(5.00, 'monthly'); + + expect(tracker.getCurrentSpending('daily')).toBe(1.00); + expect(tracker.getCurrentSpending('monthly')).toBe(5.00); + }); + }); + + // ========================================================================== + // Warning Threshold Tests + // ========================================================================== + + describe('getWarningThresholds', () => { + it('should return default thresholds', () => { + const thresholds = tracker.getWarningThresholds(); + expect(thresholds).toEqual([0.5, 0.8]); + }); + + it('should return custom thresholds from config', () => { + configManager.setBudgetConfig({ warningThresholds: [0.25, 0.5, 0.75] }); + + const thresholds = tracker.getWarningThresholds(); + expect(thresholds).toEqual([0.25, 0.5, 0.75]); + }); + }); + + describe('checkWarningThresholds', () => { + beforeEach(() => { + configManager.setBudgetConfig({ monthlyLimit: 10.00 }); + }); + + it('should return null when under first threshold', () => { + tracker.recordCost(2.00, 'monthly'); // 20% + + const warning = tracker.checkWarningThresholds(); + expect(warning).toBeNull(); + }); + + it('should return 50% warning when crossed', () => { + tracker.recordCost(5.50, 'monthly'); // 55% + + const warning = tracker.checkWarningThresholds(); + expect(warning).toBeDefined(); + expect(warning?.threshold).toBe(0.5); + expect(warning?.percentUsed).toBeCloseTo(0.55, 2); + }); + + it('should return 80% warning when crossed', () => { + tracker.recordCost(8.50, 'monthly'); // 85% + + const warning = tracker.checkWarningThresholds(); + expect(warning).toBeDefined(); + expect(warning?.threshold).toBe(0.8); + }); + + it('should return null when no monthly limit set', () => { + configManager.setBudgetConfig({ monthlyLimit: undefined }); + tracker.recordCost(100.00, 'monthly'); + + const warning = tracker.checkWarningThresholds(); + expect(warning).toBeNull(); + }); + }); + + // ========================================================================== + // Per-Session Budget Tests + // ========================================================================== + + describe('getPerSessionDefault', () => { + it('should return undefined when not configured', () => { + const defaultBudget = tracker.getPerSessionDefault(); + expect(defaultBudget).toBeUndefined(); + }); + + it('should return configured per-session default', () => { + configManager.setBudgetConfig({ perSessionDefault: 2.50 }); + + const defaultBudget = tracker.getPerSessionDefault(); + expect(defaultBudget).toBe(2.50); + }); + }); + + describe('validateSessionBudget', () => { + it('should accept any budget when no limits configured', () => { + const result = tracker.validateSessionBudget(100.00); + expect(result.valid).toBe(true); + }); + + it('should reject negative budget', () => { + const result = tracker.validateSessionBudget(-5.00); + expect(result.valid).toBe(false); + expect(result.message).toContain('positive'); + }); + + it('should warn when session budget exceeds remaining monthly budget', () => { + configManager.setBudgetConfig({ monthlyLimit: 10.00 }); + tracker.recordCost(8.00, 'monthly'); + + const result = tracker.validateSessionBudget(5.00); // Would exceed + expect(result.valid).toBe(true); // Still valid, but... + expect(result.warning).toContain('exceed'); + }); + }); +}); diff --git a/tests/config.test.ts b/tests/config.test.ts new file mode 100644 index 0000000..2325553 --- /dev/null +++ b/tests/config.test.ts @@ -0,0 +1,426 @@ +/** + * Tests for ConfigManager class + * + * TDD: These tests define the expected behavior of the ConfigManager + * before implementation. + */ + +import { describe, it, expect, beforeEach, afterEach } from 'vitest'; +import * as fs from 'fs'; +import * as path from 'path'; +import * as os from 'os'; +import { ConfigManager, DEFAULT_BUDGET_CONFIG } from '../src/config.js'; +import type { BudgetConfig } from '../src/types.js'; + +// Test fixtures directory - unique per process to avoid conflicts +const TEST_DIR = path.join(os.tmpdir(), 'parallel-cc-config-test-' + process.pid); +const TEST_CONFIG_PATH = path.join(TEST_DIR, '.parallel-cc', 'config.json'); + +describe('ConfigManager', () => { + let configManager: ConfigManager; + + beforeEach(() => { + // Create test directory + fs.mkdirSync(TEST_DIR, { recursive: true }); + }); + + afterEach(() => { + // Clean up test directory + fs.rmSync(TEST_DIR, { recursive: true, force: true }); + }); + + // ========================================================================== + // Constructor and Initialization Tests + // ========================================================================== + + describe('constructor', () => { + it('should create config file directory if it does not exist', () => { + const configPath = path.join(TEST_DIR, 'nested', 'deep', 'config.json'); + configManager = new ConfigManager(configPath); + + expect(fs.existsSync(path.dirname(configPath))).toBe(true); + }); + + it('should load existing config file', () => { + // Create a config file first + const configDir = path.dirname(TEST_CONFIG_PATH); + fs.mkdirSync(configDir, { recursive: true }); + fs.writeFileSync(TEST_CONFIG_PATH, JSON.stringify({ + budget: { + monthlyLimit: 10.00, + perSessionDefault: 1.00 + } + })); + + configManager = new ConfigManager(TEST_CONFIG_PATH); + + expect(configManager.get('budget.monthlyLimit')).toBe(10.00); + expect(configManager.get('budget.perSessionDefault')).toBe(1.00); + }); + + it('should create default config if file does not exist', () => { + configManager = new ConfigManager(TEST_CONFIG_PATH); + + expect(configManager.getAll()).toEqual({ + budget: DEFAULT_BUDGET_CONFIG + }); + }); + + it('should handle invalid JSON gracefully', () => { + // Create an invalid JSON file + const configDir = path.dirname(TEST_CONFIG_PATH); + fs.mkdirSync(configDir, { recursive: true }); + fs.writeFileSync(TEST_CONFIG_PATH, 'not valid json {{{'); + + // Should not throw, should use defaults + configManager = new ConfigManager(TEST_CONFIG_PATH); + + expect(configManager.getAll()).toEqual({ + budget: DEFAULT_BUDGET_CONFIG + }); + }); + }); + + // ========================================================================== + // get() Tests + // ========================================================================== + + describe('get', () => { + beforeEach(() => { + configManager = new ConfigManager(TEST_CONFIG_PATH); + }); + + it('should get top-level config value', () => { + configManager.set('testKey', 'testValue'); + + expect(configManager.get('testKey')).toBe('testValue'); + }); + + it('should get nested config value with dot notation', () => { + const value = configManager.get('budget.warningThresholds'); + + expect(value).toEqual(DEFAULT_BUDGET_CONFIG.warningThresholds); + }); + + it('should return undefined for non-existent key', () => { + expect(configManager.get('nonexistent')).toBeUndefined(); + expect(configManager.get('budget.nonexistent')).toBeUndefined(); + expect(configManager.get('a.b.c.d')).toBeUndefined(); + }); + + it('should return entire nested object when accessing parent key', () => { + const budget = configManager.get('budget'); + + expect(budget).toEqual(DEFAULT_BUDGET_CONFIG); + }); + }); + + // ========================================================================== + // set() Tests + // ========================================================================== + + describe('set', () => { + beforeEach(() => { + configManager = new ConfigManager(TEST_CONFIG_PATH); + }); + + it('should set top-level config value', () => { + configManager.set('newKey', 'newValue'); + + expect(configManager.get('newKey')).toBe('newValue'); + }); + + it('should set nested config value with dot notation', () => { + configManager.set('budget.monthlyLimit', 25.00); + + expect(configManager.get('budget.monthlyLimit')).toBe(25.00); + }); + + it('should create intermediate objects for deep paths', () => { + configManager.set('deep.nested.value', 42); + + expect(configManager.get('deep.nested.value')).toBe(42); + expect(configManager.get('deep.nested')).toEqual({ value: 42 }); + }); + + it('should persist changes to file', () => { + configManager.set('budget.monthlyLimit', 50.00); + + // Read directly from file + const fileContent = JSON.parse(fs.readFileSync(TEST_CONFIG_PATH, 'utf-8')); + + expect(fileContent.budget.monthlyLimit).toBe(50.00); + }); + + it('should overwrite existing values', () => { + configManager.set('budget.monthlyLimit', 10.00); + configManager.set('budget.monthlyLimit', 20.00); + + expect(configManager.get('budget.monthlyLimit')).toBe(20.00); + }); + + it('should handle null values', () => { + configManager.set('budget.monthlyLimit', null); + + expect(configManager.get('budget.monthlyLimit')).toBeNull(); + }); + + it('should handle array values', () => { + configManager.set('budget.warningThresholds', [0.3, 0.6, 0.9]); + + expect(configManager.get('budget.warningThresholds')).toEqual([0.3, 0.6, 0.9]); + }); + }); + + // ========================================================================== + // getAll() Tests + // ========================================================================== + + describe('getAll', () => { + it('should return entire config object', () => { + configManager = new ConfigManager(TEST_CONFIG_PATH); + configManager.set('budget.monthlyLimit', 15.00); + configManager.set('customKey', 'customValue'); + + const config = configManager.getAll(); + + expect(config).toHaveProperty('budget'); + expect(config).toHaveProperty('customKey'); + expect(config.budget.monthlyLimit).toBe(15.00); + expect(config.customKey).toBe('customValue'); + }); + + it('should return a copy, not a reference', () => { + configManager = new ConfigManager(TEST_CONFIG_PATH); + const config1 = configManager.getAll(); + + // Mutate the returned object + config1.budget.monthlyLimit = 999.99; + + // Original should be unchanged + expect(configManager.get('budget.monthlyLimit')).not.toBe(999.99); + }); + }); + + // ========================================================================== + // getBudgetConfig() Tests + // ========================================================================== + + describe('getBudgetConfig', () => { + beforeEach(() => { + configManager = new ConfigManager(TEST_CONFIG_PATH); + }); + + it('should return default budget config', () => { + const budgetConfig = configManager.getBudgetConfig(); + + expect(budgetConfig).toEqual(DEFAULT_BUDGET_CONFIG); + }); + + it('should return updated budget config after set', () => { + configManager.set('budget.monthlyLimit', 100.00); + configManager.set('budget.perSessionDefault', 5.00); + + const budgetConfig = configManager.getBudgetConfig(); + + expect(budgetConfig.monthlyLimit).toBe(100.00); + expect(budgetConfig.perSessionDefault).toBe(5.00); + }); + + it('should include warning thresholds', () => { + const budgetConfig = configManager.getBudgetConfig(); + + expect(budgetConfig.warningThresholds).toEqual([0.5, 0.8]); + }); + }); + + // ========================================================================== + // setBudgetConfig() Tests + // ========================================================================== + + describe('setBudgetConfig', () => { + beforeEach(() => { + configManager = new ConfigManager(TEST_CONFIG_PATH); + }); + + it('should update budget config partially', () => { + configManager.setBudgetConfig({ monthlyLimit: 50.00 }); + + const budgetConfig = configManager.getBudgetConfig(); + + expect(budgetConfig.monthlyLimit).toBe(50.00); + expect(budgetConfig.perSessionDefault).toEqual(DEFAULT_BUDGET_CONFIG.perSessionDefault); + expect(budgetConfig.warningThresholds).toEqual(DEFAULT_BUDGET_CONFIG.warningThresholds); + }); + + it('should update multiple budget config fields', () => { + configManager.setBudgetConfig({ + monthlyLimit: 75.00, + perSessionDefault: 2.50, + warningThresholds: [0.25, 0.5, 0.75] + }); + + const budgetConfig = configManager.getBudgetConfig(); + + expect(budgetConfig.monthlyLimit).toBe(75.00); + expect(budgetConfig.perSessionDefault).toBe(2.50); + expect(budgetConfig.warningThresholds).toEqual([0.25, 0.5, 0.75]); + }); + + it('should persist budget config to file', () => { + configManager.setBudgetConfig({ monthlyLimit: 30.00 }); + + // Read directly from file + const fileContent = JSON.parse(fs.readFileSync(TEST_CONFIG_PATH, 'utf-8')); + + expect(fileContent.budget.monthlyLimit).toBe(30.00); + }); + + it('should clear budget config values with undefined', () => { + configManager.setBudgetConfig({ monthlyLimit: 100.00 }); + configManager.setBudgetConfig({ monthlyLimit: undefined }); + + const budgetConfig = configManager.getBudgetConfig(); + + expect(budgetConfig.monthlyLimit).toBeUndefined(); + }); + }); + + // ========================================================================== + // delete() Tests + // ========================================================================== + + describe('delete', () => { + beforeEach(() => { + configManager = new ConfigManager(TEST_CONFIG_PATH); + }); + + it('should delete a config key', () => { + configManager.set('toDelete', 'value'); + + expect(configManager.get('toDelete')).toBe('value'); + + configManager.delete('toDelete'); + + expect(configManager.get('toDelete')).toBeUndefined(); + }); + + it('should delete nested config key', () => { + configManager.set('budget.monthlyLimit', 100.00); + + configManager.delete('budget.monthlyLimit'); + + expect(configManager.get('budget.monthlyLimit')).toBeUndefined(); + // Parent should still exist + expect(configManager.get('budget')).toBeDefined(); + }); + + it('should handle deleting non-existent key gracefully', () => { + // Should not throw + expect(() => configManager.delete('nonexistent')).not.toThrow(); + expect(() => configManager.delete('a.b.c')).not.toThrow(); + }); + + it('should persist deletion to file', () => { + configManager.set('deleteMe', 'value'); + configManager.delete('deleteMe'); + + // Read directly from file + const fileContent = JSON.parse(fs.readFileSync(TEST_CONFIG_PATH, 'utf-8')); + + expect(fileContent.deleteMe).toBeUndefined(); + }); + }); + + // ========================================================================== + // Validation Tests + // ========================================================================== + + describe('validation', () => { + beforeEach(() => { + configManager = new ConfigManager(TEST_CONFIG_PATH); + }); + + it('should reject negative budget limits', () => { + expect(() => configManager.setBudgetConfig({ monthlyLimit: -10.00 })) + .toThrow('Budget limit must be a positive number'); + }); + + it('should reject negative per-session defaults', () => { + expect(() => configManager.setBudgetConfig({ perSessionDefault: -5.00 })) + .toThrow('Per-session budget must be a positive number'); + }); + + it('should reject invalid warning thresholds', () => { + expect(() => configManager.setBudgetConfig({ warningThresholds: [-0.1, 0.5] })) + .toThrow('Warning thresholds must be between 0 and 1'); + + expect(() => configManager.setBudgetConfig({ warningThresholds: [0.5, 1.5] })) + .toThrow('Warning thresholds must be between 0 and 1'); + }); + + it('should accept valid budget config', () => { + expect(() => configManager.setBudgetConfig({ + monthlyLimit: 100.00, + perSessionDefault: 5.00, + warningThresholds: [0.25, 0.5, 0.75] + })).not.toThrow(); + }); + + it('should accept zero as valid budget limit', () => { + // Zero means "no limit" or "disabled" + expect(() => configManager.setBudgetConfig({ monthlyLimit: 0 })).not.toThrow(); + }); + }); + + // ========================================================================== + // Persistence and Concurrency Tests + // ========================================================================== + + describe('persistence', () => { + it('should preserve config across instances', () => { + const configManager1 = new ConfigManager(TEST_CONFIG_PATH); + configManager1.set('persistent', 'value'); + + const configManager2 = new ConfigManager(TEST_CONFIG_PATH); + + expect(configManager2.get('persistent')).toBe('value'); + }); + + it('should write valid JSON', () => { + configManager = new ConfigManager(TEST_CONFIG_PATH); + configManager.set('complex', { + nested: { deep: { value: 'test' } }, + array: [1, 2, 3] + }); + + // Should not throw when parsing + const content = fs.readFileSync(TEST_CONFIG_PATH, 'utf-8'); + expect(() => JSON.parse(content)).not.toThrow(); + }); + + it('should format JSON with indentation for readability', () => { + configManager = new ConfigManager(TEST_CONFIG_PATH); + configManager.set('test', 'value'); + + const content = fs.readFileSync(TEST_CONFIG_PATH, 'utf-8'); + + // Should contain newlines (formatted) + expect(content).toContain('\n'); + }); + }); + + // ========================================================================== + // Default Config Tests + // ========================================================================== + + describe('DEFAULT_BUDGET_CONFIG', () => { + it('should have expected default values', () => { + expect(DEFAULT_BUDGET_CONFIG).toEqual({ + monthlyLimit: undefined, + perSessionDefault: undefined, + warningThresholds: [0.5, 0.8] + }); + }); + }); +}); diff --git a/tests/e2b/budget-enforcement.test.ts b/tests/e2b/budget-enforcement.test.ts new file mode 100644 index 0000000..7219ad3 --- /dev/null +++ b/tests/e2b/budget-enforcement.test.ts @@ -0,0 +1,382 @@ +/** + * Tests for budget enforcement in SandboxManager + * + * Tests budget limits during sandbox execution including: + * - Setting budget limits per sandbox + * - Budget warnings at configurable thresholds + * - Hard termination when budget exceeded + * - Integration with cost calculations + */ + +import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest'; +import { SandboxManager } from '../../src/e2b/sandbox-manager.js'; +import { SandboxStatus, BudgetExceededError } from '../../src/types.js'; +import type { Logger } from '../../src/logger.js'; +import type { BudgetWarning } from '../../src/types.js'; + +// Mock E2B SDK +vi.mock('e2b', () => ({ + Sandbox: { + create: vi.fn(), + connect: vi.fn() + } +})); + +// Create mock logger +const createMockLogger = (): Logger => ({ + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + debug: vi.fn() +}); + +describe('SandboxManager Budget Enforcement', () => { + let manager: SandboxManager; + let mockLogger: Logger; + let mockSandbox: any; + + beforeEach(async () => { + vi.clearAllMocks(); + vi.useFakeTimers(); + + // Create mock logger + mockLogger = createMockLogger(); + + // Create mock sandbox instance + mockSandbox = { + sandboxId: 'test-sandbox-budget', + close: vi.fn().mockResolvedValue(undefined), + kill: vi.fn().mockResolvedValue(undefined), + isRunning: vi.fn().mockResolvedValue(true), + setTimeout: vi.fn().mockResolvedValue(undefined), + metadata: {} + }; + + // Setup E2B SDK mock + const { Sandbox } = await import('e2b'); + vi.mocked(Sandbox.create).mockResolvedValue(mockSandbox); + + // Set environment variable for API key + process.env.E2B_API_KEY = 'test-api-key-12345'; + + // Create manager instance with 60 minute timeout and custom budget warning thresholds + manager = new SandboxManager(mockLogger, { + timeoutMinutes: 60, + warningThresholds: [30, 50] // 30min and 50min warnings + }); + }); + + afterEach(() => { + vi.clearAllTimers(); + vi.useRealTimers(); + delete process.env.E2B_API_KEY; + }); + + // ========================================================================== + // Budget Limit Setting Tests + // ========================================================================== + + describe('setBudgetLimit', () => { + it('should set budget limit for sandbox', async () => { + const { sandboxId } = await manager.createSandbox('session-1'); + + manager.setBudgetLimit(sandboxId, 1.00); + + const limit = manager.getBudgetLimit(sandboxId); + expect(limit).toBe(1.00); + }); + + it('should allow updating budget limit', async () => { + const { sandboxId } = await manager.createSandbox('session-1'); + + manager.setBudgetLimit(sandboxId, 1.00); + manager.setBudgetLimit(sandboxId, 2.00); + + const limit = manager.getBudgetLimit(sandboxId); + expect(limit).toBe(2.00); + }); + + it('should return undefined for sandbox without budget limit', async () => { + const { sandboxId } = await manager.createSandbox('session-1'); + + const limit = manager.getBudgetLimit(sandboxId); + expect(limit).toBeUndefined(); + }); + + it('should reject negative budget limits', async () => { + const { sandboxId } = await manager.createSandbox('session-1'); + + expect(() => manager.setBudgetLimit(sandboxId, -1.00)) + .toThrow('Budget limit must be a positive number'); + }); + + it('should accept zero as "no limit"', async () => { + const { sandboxId } = await manager.createSandbox('session-1'); + + manager.setBudgetLimit(sandboxId, 0); + + const limit = manager.getBudgetLimit(sandboxId); + expect(limit).toBe(0); + }); + }); + + // ========================================================================== + // Budget Check Tests + // ========================================================================== + + describe('checkBudgetLimit', () => { + it('should return null when no budget limit set', async () => { + const { sandboxId } = await manager.createSandbox('session-1'); + + // Simulate time passing + vi.advanceTimersByTime(30 * 60 * 1000); // 30 minutes + + const warning = await manager.checkBudgetLimit(sandboxId); + expect(warning).toBeNull(); + }); + + it('should return null when under budget', async () => { + const { sandboxId } = await manager.createSandbox('session-1'); + manager.setBudgetLimit(sandboxId, 1.00); // $1.00 limit + + // 10 minutes = ~$0.017 (well under $1.00) + vi.advanceTimersByTime(10 * 60 * 1000); + + const warning = await manager.checkBudgetLimit(sandboxId); + expect(warning).toBeNull(); + }); + + it('should return soft warning at 50% threshold', async () => { + const { sandboxId } = await manager.createSandbox('session-1'); + manager.setBudgetLimit(sandboxId, 0.10); // $0.10 limit + + // 30 minutes = ~$0.05, which is 50% of $0.10 + vi.advanceTimersByTime(30 * 60 * 1000); + + const warning = await manager.checkBudgetLimit(sandboxId); + + expect(warning).not.toBeNull(); + expect(warning?.warningLevel).toBe('soft'); + expect(warning?.percentUsed).toBeGreaterThanOrEqual(0.5); + }); + + it('should return soft warning at 80% threshold', async () => { + const { sandboxId } = await manager.createSandbox('session-1'); + manager.setBudgetLimit(sandboxId, 0.10); // $0.10 limit + + // 49 minutes = ~$0.0817, which is >80% of $0.10 + vi.advanceTimersByTime(49 * 60 * 1000); + + const warning = await manager.checkBudgetLimit(sandboxId); + + expect(warning).not.toBeNull(); + expect(warning?.warningLevel).toBe('soft'); + // Use toBeCloseTo for floating point comparison + expect(warning?.percentUsed).toBeCloseTo(0.817, 2); + }); + + it('should throw BudgetExceededError when budget exceeded', async () => { + const { sandboxId } = await manager.createSandbox('session-1'); + manager.setBudgetLimit(sandboxId, 0.05); // $0.05 limit (30 minutes) + + // 35 minutes = ~$0.058, exceeds $0.05 + vi.advanceTimersByTime(35 * 60 * 1000); + + await expect(manager.checkBudgetLimit(sandboxId)) + .rejects.toThrow(BudgetExceededError); + }); + + it('should terminate sandbox when budget exceeded', async () => { + const { sandboxId } = await manager.createSandbox('session-1'); + manager.setBudgetLimit(sandboxId, 0.05); + + vi.advanceTimersByTime(35 * 60 * 1000); + + try { + await manager.checkBudgetLimit(sandboxId); + } catch (error) { + // Expected + } + + expect(mockSandbox.kill).toHaveBeenCalled(); + }); + + it('should not issue duplicate warnings', async () => { + const { sandboxId } = await manager.createSandbox('session-1'); + manager.setBudgetLimit(sandboxId, 0.10); + + // First check at 50% + vi.advanceTimersByTime(30 * 60 * 1000); + const warning1 = await manager.checkBudgetLimit(sandboxId); + expect(warning1).not.toBeNull(); + + // Second check still at 50% - should not warn again + vi.advanceTimersByTime(1 * 60 * 1000); + const warning2 = await manager.checkBudgetLimit(sandboxId); + expect(warning2).toBeNull(); + }); + }); + + // ========================================================================== + // Budget Warning Details Tests + // ========================================================================== + + describe('budget warning details', () => { + it('should include current cost in warning', async () => { + const { sandboxId } = await manager.createSandbox('session-1'); + manager.setBudgetLimit(sandboxId, 0.10); + + vi.advanceTimersByTime(30 * 60 * 1000); + + const warning = await manager.checkBudgetLimit(sandboxId); + + expect(warning?.currentCost).toBeCloseTo(0.05, 2); + }); + + it('should include budget limit in warning', async () => { + const { sandboxId } = await manager.createSandbox('session-1'); + manager.setBudgetLimit(sandboxId, 0.10); + + vi.advanceTimersByTime(30 * 60 * 1000); + + const warning = await manager.checkBudgetLimit(sandboxId); + + expect(warning?.budgetLimit).toBe(0.10); + }); + + it('should include percent used in warning', async () => { + const { sandboxId } = await manager.createSandbox('session-1'); + manager.setBudgetLimit(sandboxId, 0.10); + + vi.advanceTimersByTime(30 * 60 * 1000); + + const warning = await manager.checkBudgetLimit(sandboxId); + + expect(warning?.percentUsed).toBeGreaterThanOrEqual(0.5); + expect(warning?.percentUsed).toBeLessThanOrEqual(1.0); + }); + + it('should include descriptive message in warning', async () => { + const { sandboxId } = await manager.createSandbox('session-1'); + manager.setBudgetLimit(sandboxId, 0.10); + + vi.advanceTimersByTime(30 * 60 * 1000); + + const warning = await manager.checkBudgetLimit(sandboxId); + + expect(warning?.message).toContain('budget'); + expect(warning?.message).toContain('$'); + }); + }); + + // ========================================================================== + // Cost Calculation Integration Tests + // ========================================================================== + + describe('cost calculation integration', () => { + it('should calculate cost based on elapsed time', async () => { + const { sandboxId } = await manager.createSandbox('session-1'); + + vi.advanceTimersByTime(60 * 60 * 1000); // 60 minutes + + const cost = manager.getEstimatedCost(sandboxId); + + // $0.10/hour = $0.10 for 60 minutes + expect(cost).toBe('$0.10'); + }); + + it('should return null for unknown sandbox', () => { + const cost = manager.getEstimatedCost('unknown-sandbox'); + expect(cost).toBeNull(); + }); + + it('should calculate cost correctly for partial hours', async () => { + const { sandboxId } = await manager.createSandbox('session-1'); + + vi.advanceTimersByTime(30 * 60 * 1000); // 30 minutes + + const cost = manager.getEstimatedCost(sandboxId); + + // $0.10/hour = $0.05 for 30 minutes + expect(cost).toBe('$0.05'); + }); + }); + + // ========================================================================== + // Budget Thresholds Configuration Tests + // ========================================================================== + + describe('budget thresholds configuration', () => { + it('should use default thresholds (50%, 80%)', async () => { + const defaultManager = new SandboxManager(mockLogger); + const { sandboxId } = await defaultManager.createSandbox('session-1'); + defaultManager.setBudgetLimit(sandboxId, 0.10); + + // At 50% + vi.advanceTimersByTime(30 * 60 * 1000); + const warning = await defaultManager.checkBudgetLimit(sandboxId); + + expect(warning?.percentUsed).toBeGreaterThanOrEqual(0.5); + }); + + it('should support custom thresholds', async () => { + const customManager = new SandboxManager(mockLogger, { + budgetWarningThresholds: [0.25, 0.5, 0.75] + }); + const { sandboxId } = await customManager.createSandbox('session-1'); + customManager.setBudgetLimit(sandboxId, 0.10); + + // At 25% + vi.advanceTimersByTime(15 * 60 * 1000); // ~$0.025 + const warning = await customManager.checkBudgetLimit(sandboxId); + + expect(warning).not.toBeNull(); + expect(warning?.percentUsed).toBeGreaterThanOrEqual(0.25); + }); + }); + + // ========================================================================== + // Edge Cases Tests + // ========================================================================== + + describe('edge cases', () => { + it('should handle sandbox not found gracefully', async () => { + // For non-existent sandbox with no budget set, should return null + const warning = await manager.checkBudgetLimit('non-existent-sandbox'); + expect(warning).toBeNull(); + // No warning logged because no budget was set for this sandbox + }); + + it('should handle zero budget limit (no spending allowed)', async () => { + const { sandboxId } = await manager.createSandbox('session-1'); + manager.setBudgetLimit(sandboxId, 0); + + // Any time = exceeded (since budget is 0) + vi.advanceTimersByTime(1 * 60 * 1000); + + // Zero budget means "disabled" - should not trigger + const warning = await manager.checkBudgetLimit(sandboxId); + expect(warning).toBeNull(); + }); + + it('should clean up budget tracking when sandbox terminated', async () => { + const { sandboxId } = await manager.createSandbox('session-1'); + manager.setBudgetLimit(sandboxId, 1.00); + + await manager.terminateSandbox(sandboxId); + + const limit = manager.getBudgetLimit(sandboxId); + expect(limit).toBeUndefined(); + }); + + it('should preserve budget limit after health check', async () => { + const { sandboxId } = await manager.createSandbox('session-1'); + manager.setBudgetLimit(sandboxId, 0.50); + + await manager.monitorSandboxHealth(sandboxId, false); + + const limit = manager.getBudgetLimit(sandboxId); + expect(limit).toBe(0.50); + }); + }); +}); From 58c778db2d607d7e2b50e80619f4df170c42d9dc Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 29 Jan 2026 23:27:54 +0000 Subject: [PATCH 2/7] fix(budget): address precision and reliability issues in budget tracking Fixes the following issues identified in code review: 1. Float precision bug: Convert warning threshold Set keys to integer percentages (0-100) to avoid floating-point comparison issues with values like 0.8 2. Hard-coded pricing: Move E2B hourly rate ($0.10/hour) to configurable e2bHourlyRate in BudgetConfig, allowing updates when pricing changes 3. Synchronous I/O: Add debounced writes to ConfigManager (100ms window) with flushSync() for immediate persistence and cancelPendingWrites() for clean test teardown 4. UTC documentation: Add comprehensive JSDoc explaining that period calculations use UTC dates, with timezone implications noted 5. Warning reset: Add automatic warning reset when budget period rolls over (e.g., new month starts) https://claude.ai/code/session_01CDYWMZJdg8FH55bgQzNmds --- src/budget-tracker.ts | 38 +++++++++++++++++++++-- src/config.ts | 59 +++++++++++++++++++++++++++++++++--- src/e2b/sandbox-manager.ts | 25 +++++++++------ src/types.ts | 2 ++ tests/budget-tracker.test.ts | 2 ++ tests/config.test.ts | 23 +++++++++++++- 6 files changed, 132 insertions(+), 17 deletions(-) diff --git a/src/budget-tracker.ts b/src/budget-tracker.ts index 655e70b..52e8ea5 100644 --- a/src/budget-tracker.ts +++ b/src/budget-tracker.ts @@ -52,11 +52,17 @@ export interface SessionBudgetValidation { * - Budget limit enforcement with configurable thresholds * - Session cost recording * - Budget status reports + * + * Note: Uses integer percentages (0-100) internally for Set keys + * to avoid floating-point precision issues. */ export class BudgetTracker { private db: SessionDB; private configManager: ConfigManager; + // Use integer percentages (0-100) as Set keys to avoid float precision issues private warningsIssued: Set = new Set(); + // Track current period to reset warnings on period rollover + private currentPeriodKey: string | null = null; constructor(db: SessionDB, configManager: ConfigManager) { this.db = db; @@ -66,9 +72,19 @@ export class BudgetTracker { /** * Calculate the start date for a given period * + * **IMPORTANT: UTC-based calculation** + * All period calculations use UTC dates to ensure consistency across timezones. + * This means: + * - Daily periods reset at midnight UTC (not local time) + * - Weekly periods start on Monday UTC + * - Monthly periods start on the 1st UTC + * + * Users in different timezones may see period boundaries at different local times. + * For example, a user in PST (-8) will see daily periods reset at 4pm local time. + * * @param period - Budget period type * @param date - Reference date (default: now) - * @returns ISO date string (YYYY-MM-DD) + * @returns ISO date string (YYYY-MM-DD) in UTC */ getPeriodStart(period: BudgetPeriod, date: Date = new Date()): string { const year = date.getUTCFullYear(); @@ -99,6 +115,9 @@ export class BudgetTracker { /** * Get or create a budget tracking record for the current period * + * Automatically resets warning tracking when the period rolls over + * to a new time window. + * * @param period - Budget period type * @returns Budget tracking record */ @@ -106,6 +125,14 @@ export class BudgetTracker { const periodStart = this.getPeriodStart(period); const budgetConfig = this.configManager.getBudgetConfig(); + // Check for period rollover and reset warnings if needed + const periodKey = `${period}:${periodStart}`; + if (this.currentPeriodKey !== null && this.currentPeriodKey !== periodKey) { + // Period has rolled over - reset warning tracking + this.warningsIssued.clear(); + } + this.currentPeriodKey = periodKey; + // Get budget limit for this period type let budgetLimit: number | undefined; if (period === 'monthly' && budgetConfig.monthlyLimit !== undefined) { @@ -258,6 +285,9 @@ export class BudgetTracker { /** * Check if any warning thresholds have been crossed * + * Uses integer percentages (0-100) for Set keys to avoid + * floating-point precision issues with values like 0.8. + * * @returns Warning info if threshold crossed, null otherwise */ checkWarningThresholds(): BudgetThresholdWarning | null { @@ -275,8 +305,10 @@ export class BudgetTracker { // Find the highest threshold that has been crossed but not yet warned about for (const threshold of thresholds) { - if (percentUsed >= threshold && !this.warningsIssued.has(threshold)) { - this.warningsIssued.add(threshold); + // Convert to integer percentage (0-100) for reliable Set lookup + const thresholdInt = Math.round(threshold * 100); + if (percentUsed >= threshold && !this.warningsIssued.has(thresholdInt)) { + this.warningsIssued.add(thresholdInt); return { threshold, diff --git a/src/config.ts b/src/config.ts index 9a3ab74..4e6a1bc 100644 --- a/src/config.ts +++ b/src/config.ts @@ -15,7 +15,8 @@ import type { BudgetConfig } from './types.js'; export const DEFAULT_BUDGET_CONFIG: BudgetConfig = { monthlyLimit: undefined, perSessionDefault: undefined, - warningThresholds: [0.5, 0.8] + warningThresholds: [0.5, 0.8], + e2bHourlyRate: 0.10 // Default E2B pricing: $0.10/hour }; /** @@ -42,11 +43,14 @@ export const DEFAULT_CONFIG_PATH = pathJoin(homedir(), '.parallel-cc', 'config.j * - JSON file storage with automatic directory creation * - Dot notation support for nested keys (e.g., "budget.monthlyLimit") * - Validation for budget-related settings - * - Automatic persistence on changes + * - Debounced writes to avoid excessive disk I/O during rapid changes */ export class ConfigManager { private configPath: string; private config: Config; + private saveTimer: ReturnType | null = null; + private readonly debounceMs: number = 100; // 100ms debounce window + private pendingSave: boolean = false; /** * Create a new ConfigManager @@ -109,11 +113,58 @@ export class ConfigManager { } /** - * Save config to file + * Save config to file (debounced to avoid excessive I/O) + * + * Multiple rapid calls will be batched into a single write. + * Use flushSync() to force immediate write when needed. */ private save(config?: Config): void { const toSave = config ?? this.config; - writeFileSync(this.configPath, JSON.stringify(toSave, null, 2)); + this.config = toSave; + this.pendingSave = true; + + // Clear existing timer if any + if (this.saveTimer) { + clearTimeout(this.saveTimer); + } + + // Schedule debounced write + this.saveTimer = setTimeout(() => { + this.flushSync(); + }, this.debounceMs); + } + + /** + * Immediately flush pending config changes to disk (synchronous) + * + * Call this when you need to ensure config is persisted immediately, + * such as before process exit. + */ + flushSync(): void { + if (this.saveTimer) { + clearTimeout(this.saveTimer); + this.saveTimer = null; + } + + if (this.pendingSave) { + // Only write if directory still exists (handles test cleanup race conditions) + const dir = dirname(this.configPath); + if (existsSync(dir)) { + writeFileSync(this.configPath, JSON.stringify(this.config, null, 2)); + } + this.pendingSave = false; + } + } + + /** + * Cancel any pending writes (for cleanup/testing) + */ + cancelPendingWrites(): void { + if (this.saveTimer) { + clearTimeout(this.saveTimer); + this.saveTimer = null; + } + this.pendingSave = false; } /** diff --git a/src/e2b/sandbox-manager.ts b/src/e2b/sandbox-manager.ts index 1bf6f16..64e5a8e 100644 --- a/src/e2b/sandbox-manager.ts +++ b/src/e2b/sandbox-manager.ts @@ -31,9 +31,11 @@ export interface TemplateApplicationResult { error?: string; } -// Extended config type with budget thresholds +// Extended config type with budget thresholds and pricing interface ExtendedE2BSessionConfig extends E2BSessionConfig { budgetWarningThresholds?: number[]; + /** E2B hourly rate in USD - allows updating if pricing changes */ + e2bHourlyRate?: number; } // Default configuration @@ -44,7 +46,8 @@ const DEFAULT_CONFIG: Required = { sandboxImage: (envTemplate && envTemplate.length > 0) ? envTemplate : 'anthropic-claude-code', // E2B template with pre-installed Claude Code timeoutMinutes: 60, warningThresholds: [30, 50], - budgetWarningThresholds: [0.5, 0.8] // Default: warn at 50% and 80% of budget + budgetWarningThresholds: [0.5, 0.8], // Default: warn at 50% and 80% of budget + e2bHourlyRate: 0.10 // Default E2B pricing: $0.10/hour (configurable if pricing changes) }; // Security constants @@ -425,11 +428,14 @@ export class SandboxManager { } // Check soft warning thresholds + // Uses integer percentages (0-100) for Set keys to avoid float precision issues const thresholds = this.config.budgetWarningThresholds.sort((a, b) => b - a); // Sort descending for (const threshold of thresholds) { - if (percentUsed >= threshold && !warningsIssued.has(threshold)) { - warningsIssued.add(threshold); + // Convert to integer percentage (0-100) for reliable Set lookup + const thresholdInt = Math.round(threshold * 100); + if (percentUsed >= threshold && !warningsIssued.has(thresholdInt)) { + warningsIssued.add(thresholdInt); this.budgetWarningsIssued.set(sandboxId, warningsIssued); const warning: BudgetWarning = { @@ -462,12 +468,13 @@ export class SandboxManager { /** * Calculate estimated cost as a number (for budget calculations) * + * Uses configurable hourly rate to allow updating when E2B pricing changes. + * * @param elapsedMinutes - Elapsed minutes * @returns Estimated cost in USD */ private calculateEstimatedCostNumeric(elapsedMinutes: number): number { - // E2B pricing: ~$0.10/hour for basic compute - const costPerMinute = 0.10 / 60; + const costPerMinute = this.config.e2bHourlyRate / 60; return elapsedMinutes * costPerMinute; } @@ -630,13 +637,13 @@ export class SandboxManager { /** * Calculate estimated cost based on elapsed time * + * Uses configurable hourly rate to allow updating when E2B pricing changes. + * * @param elapsedMinutes - Elapsed minutes * @returns Estimated cost string (e.g., "$0.50") */ private calculateEstimatedCost(elapsedMinutes: number): string { - // E2B pricing: ~$0.10/hour for basic compute - const costPerMinute = 0.10 / 60; - const estimatedCost = elapsedMinutes * costPerMinute; + const estimatedCost = this.calculateEstimatedCostNumeric(elapsedMinutes); return `$${estimatedCost.toFixed(2)}`; } diff --git a/src/types.ts b/src/types.ts index 1479fb1..9a98974 100644 --- a/src/types.ts +++ b/src/types.ts @@ -622,6 +622,8 @@ export interface BudgetConfig { perSessionDefault?: number; /** Warning thresholds as percentages (e.g., [0.5, 0.8] for 50% and 80%) */ warningThresholds?: number[]; + /** E2B hourly rate in USD (default: 0.10) - allows updating if pricing changes */ + e2bHourlyRate?: number; } /** diff --git a/tests/budget-tracker.test.ts b/tests/budget-tracker.test.ts index 1688451..a3c01a0 100644 --- a/tests/budget-tracker.test.ts +++ b/tests/budget-tracker.test.ts @@ -37,6 +37,8 @@ describe('BudgetTracker', () => { }); afterEach(() => { + // Cancel any pending debounced writes to avoid race conditions with directory cleanup + configManager.cancelPendingWrites(); // Close database and clean up db.close(); fs.rmSync(TEST_DIR, { recursive: true, force: true }); diff --git a/tests/config.test.ts b/tests/config.test.ts index 2325553..9ff6743 100644 --- a/tests/config.test.ts +++ b/tests/config.test.ts @@ -25,6 +25,10 @@ describe('ConfigManager', () => { }); afterEach(() => { + // Cancel any pending debounced writes to avoid race conditions with directory cleanup + if (configManager) { + configManager.cancelPendingWrites(); + } // Clean up test directory fs.rmSync(TEST_DIR, { recursive: true, force: true }); }); @@ -145,6 +149,8 @@ describe('ConfigManager', () => { it('should persist changes to file', () => { configManager.set('budget.monthlyLimit', 50.00); + // Flush debounced writes to ensure file is updated + configManager.flushSync(); // Read directly from file const fileContent = JSON.parse(fs.readFileSync(TEST_CONFIG_PATH, 'utf-8')); @@ -269,6 +275,8 @@ describe('ConfigManager', () => { it('should persist budget config to file', () => { configManager.setBudgetConfig({ monthlyLimit: 30.00 }); + // Flush debounced writes to ensure file is updated + configManager.flushSync(); // Read directly from file const fileContent = JSON.parse(fs.readFileSync(TEST_CONFIG_PATH, 'utf-8')); @@ -324,6 +332,8 @@ describe('ConfigManager', () => { it('should persist deletion to file', () => { configManager.set('deleteMe', 'value'); configManager.delete('deleteMe'); + // Flush debounced writes to ensure file is updated + configManager.flushSync(); // Read directly from file const fileContent = JSON.parse(fs.readFileSync(TEST_CONFIG_PATH, 'utf-8')); @@ -381,10 +391,16 @@ describe('ConfigManager', () => { it('should preserve config across instances', () => { const configManager1 = new ConfigManager(TEST_CONFIG_PATH); configManager1.set('persistent', 'value'); + // Flush to ensure file is written before creating second instance + configManager1.flushSync(); const configManager2 = new ConfigManager(TEST_CONFIG_PATH); expect(configManager2.get('persistent')).toBe('value'); + + // Clean up both instances + configManager1.cancelPendingWrites(); + configManager2.cancelPendingWrites(); }); it('should write valid JSON', () => { @@ -393,6 +409,8 @@ describe('ConfigManager', () => { nested: { deep: { value: 'test' } }, array: [1, 2, 3] }); + // Flush to ensure file is written + configManager.flushSync(); // Should not throw when parsing const content = fs.readFileSync(TEST_CONFIG_PATH, 'utf-8'); @@ -402,6 +420,8 @@ describe('ConfigManager', () => { it('should format JSON with indentation for readability', () => { configManager = new ConfigManager(TEST_CONFIG_PATH); configManager.set('test', 'value'); + // Flush to ensure file is written + configManager.flushSync(); const content = fs.readFileSync(TEST_CONFIG_PATH, 'utf-8'); @@ -419,7 +439,8 @@ describe('ConfigManager', () => { expect(DEFAULT_BUDGET_CONFIG).toEqual({ monthlyLimit: undefined, perSessionDefault: undefined, - warningThresholds: [0.5, 0.8] + warningThresholds: [0.5, 0.8], + e2bHourlyRate: 0.10 }); }); }); From 4c9ced3ed056b7a6376f1f23ef690e29fcb81a25 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 29 Jan 2026 23:41:03 +0000 Subject: [PATCH 3/7] fix(budget): address critical and medium priority issues 1. Fix database connection leak in budget-status command - Changed from try/catch to try/catch/finally pattern - Database connection now closes even when errors occur - Prevents connection exhaustion from repeated failed CLI calls 2. Add key validation to ConfigManager.set() - Prevent prototype pollution attacks (__proto__, constructor, prototype) - Reject keys starting with underscore (reserved for internal use) - Validate key format (must start with letter, alphanumeric/hyphen/underscore) - Also added validation to delete() method Added 5 new tests for key validation. https://claude.ai/code/session_01CDYWMZJdg8FH55bgQzNmds --- src/cli.ts | 10 +++++++--- src/config.ts | 35 +++++++++++++++++++++++++++++++++++ tests/config.test.ts | 33 +++++++++++++++++++++++++++++++++ 3 files changed, 75 insertions(+), 3 deletions(-) diff --git a/src/cli.ts b/src/cli.ts index cd14292..d00dda0 100644 --- a/src/cli.ts +++ b/src/cli.ts @@ -3173,8 +3173,9 @@ program .option('--period ', 'Show specific period (daily, weekly, monthly)', 'monthly') .option('--json', 'Output as JSON') .action(async (options) => { + let db: SessionDB | null = null; try { - const db = new SessionDB(); + db = new SessionDB(); // Run migration if needed if (!db.hasBudgetTrackingTable()) { @@ -3228,8 +3229,6 @@ program console.log(` Remaining budget: $${status.remainingBudget.toFixed(2)}`); } } - - db.close(); } catch (error) { const errorMessage = error instanceof Error ? error.message : 'Unknown error'; if (options.json) { @@ -3238,6 +3237,11 @@ program console.error(chalk.red(`✗ Failed to get budget status: ${errorMessage}`)); } process.exit(1); + } finally { + // Always close database connection + if (db) { + db.close(); + } } }); diff --git a/src/config.ts b/src/config.ts index 4e6a1bc..39ec2a0 100644 --- a/src/config.ts +++ b/src/config.ts @@ -190,14 +190,42 @@ export class ConfigManager { return current; } + /** + * Validate that a key part is safe (prevents prototype pollution) + */ + private validateKeyPart(part: string): void { + // Prevent prototype pollution attacks + const dangerousKeys = ['__proto__', 'constructor', 'prototype']; + if (dangerousKeys.includes(part)) { + throw new Error(`Invalid config key: "${part}" is a reserved property`); + } + + // Prevent keys starting with underscore (internal properties) + if (part.startsWith('_')) { + throw new Error(`Invalid config key: keys starting with underscore are reserved`); + } + + // Validate key format (alphanumeric, hyphens, underscores only for middle parts) + if (!/^[a-zA-Z][a-zA-Z0-9_-]*$/.test(part)) { + throw new Error(`Invalid config key format: "${part}" must start with a letter and contain only alphanumeric characters, hyphens, or underscores`); + } + } + /** * Set a config value by key (supports dot notation) * * @param key - Config key (e.g., "budget.monthlyLimit") * @param value - Value to set + * @throws Error if key is invalid or reserved */ set(key: string, value: unknown): void { const parts = key.split('.'); + + // Validate all key parts + for (const part of parts) { + this.validateKeyPart(part); + } + let current: Record = this.config; // Navigate/create path to parent @@ -220,9 +248,16 @@ export class ConfigManager { * Delete a config key (supports dot notation) * * @param key - Config key to delete + * @throws Error if key is invalid or reserved */ delete(key: string): void { const parts = key.split('.'); + + // Validate all key parts + for (const part of parts) { + this.validateKeyPart(part); + } + let current: Record = this.config; // Navigate to parent diff --git a/tests/config.test.ts b/tests/config.test.ts index 9ff6743..8843d79 100644 --- a/tests/config.test.ts +++ b/tests/config.test.ts @@ -351,6 +351,39 @@ describe('ConfigManager', () => { configManager = new ConfigManager(TEST_CONFIG_PATH); }); + // Key validation tests + it('should reject prototype pollution keys', () => { + expect(() => configManager.set('__proto__', 'value')) + .toThrow('reserved property'); + expect(() => configManager.set('constructor', 'value')) + .toThrow('reserved property'); + expect(() => configManager.set('prototype', 'value')) + .toThrow('reserved property'); + }); + + it('should reject keys starting with underscore', () => { + expect(() => configManager.set('_private', 'value')) + .toThrow('keys starting with underscore are reserved'); + }); + + it('should reject keys with invalid format', () => { + expect(() => configManager.set('123invalid', 'value')) + .toThrow('must start with a letter'); + expect(() => configManager.set('has spaces', 'value')) + .toThrow('must start with a letter'); + }); + + it('should reject dangerous nested keys', () => { + expect(() => configManager.set('budget.__proto__', 'value')) + .toThrow('reserved property'); + }); + + it('should allow valid keys with hyphens and numbers', () => { + expect(() => configManager.set('valid-key', 'value')).not.toThrow(); + expect(() => configManager.set('key123', 'value')).not.toThrow(); + expect(() => configManager.set('camelCase', 'value')).not.toThrow(); + }); + it('should reject negative budget limits', () => { expect(() => configManager.setBudgetConfig({ monthlyLimit: -10.00 })) .toThrow('Budget limit must be a positive number'); From 8acc228bb7cddb94d47084fa4275cde153ff0402 Mon Sep 17 00:00:00 2001 From: Test User Date: Thu, 29 Jan 2026 22:41:28 -0700 Subject: [PATCH 4/7] fix: address code review feedback for budget tracking Fixes from PR review comments: config.ts: - Use structuredClone for DEFAULT_BUDGET_CONFIG to prevent shared refs - Add try/catch in flushSync for ENOENT handling - Add e2bHourlyRate validation in setBudgetConfig (non-negative) cli.ts: - Add NaN validation for warning threshold parsing - Add type checks before toFixed() calls (typeof === 'number') - Validate --period option against allowed values - Guard against division by zero in budget status display --- src/cli.ts | 31 ++++++++++++++++++++++++------- src/config.ts | 19 +++++++++++++++---- 2 files changed, 39 insertions(+), 11 deletions(-) diff --git a/src/cli.ts b/src/cli.ts index d00dda0..e4cb7cd 100644 --- a/src/cli.ts +++ b/src/cli.ts @@ -3048,8 +3048,12 @@ configCmd if (normalizedKey.includes('budget')) { // Budget values are numbers or arrays if (normalizedKey.includes('Thresholds')) { - // Parse as array of numbers - parsedValue = value.split(',').map(v => parseFloat(v.trim())); + // Parse as array of numbers with NaN validation + const parsed = value.split(',').map(v => parseFloat(v.trim())); + if (parsed.some(v => isNaN(v))) { + throw new Error(`Invalid threshold value in: ${value}`); + } + parsedValue = parsed; } else { // Parse as number parsedValue = parseFloat(value); @@ -3136,9 +3140,9 @@ configCmd // Display budget config console.log(chalk.cyan('Budget Settings:')); const budget = config.budget; - console.log(` monthly-limit: ${budget.monthlyLimit !== undefined ? `$${budget.monthlyLimit.toFixed(2)}` : chalk.dim('(not set)')}`); - console.log(` per-session-default: ${budget.perSessionDefault !== undefined ? `$${budget.perSessionDefault.toFixed(2)}` : chalk.dim('(not set)')}`); - console.log(` warning-thresholds: ${budget.warningThresholds ? budget.warningThresholds.map(t => `${(t * 100).toFixed(0)}%`).join(', ') : chalk.dim('(not set)')}`); + console.log(` monthly-limit: ${typeof budget.monthlyLimit === 'number' ? `$${budget.monthlyLimit.toFixed(2)}` : chalk.dim('(not set)')}`); + console.log(` per-session-default: ${typeof budget.perSessionDefault === 'number' ? `$${budget.perSessionDefault.toFixed(2)}` : chalk.dim('(not set)')}`); + console.log(` warning-thresholds: ${Array.isArray(budget.warningThresholds) ? budget.warningThresholds.map(t => `${(t * 100).toFixed(0)}%`).join(', ') : chalk.dim('(not set)')}`); // Display other config keys const otherKeys = Object.keys(config).filter(k => k !== 'budget'); @@ -3186,7 +3190,18 @@ program const configManager = new ConfigManager(); const tracker = new BudgetTracker(db, configManager); - const period = options.period as 'daily' | 'weekly' | 'monthly'; + // Validate period option + const allowedPeriods = ['daily', 'weekly', 'monthly'] as const; + if (!allowedPeriods.includes(options.period)) { + const message = `Invalid period: ${options.period}. Use daily, weekly, or monthly.`; + if (options.json) { + console.log(JSON.stringify({ error: message })); + } else { + console.error(chalk.red(`✗ ${message}`)); + } + process.exit(1); + } + const period = options.period as typeof allowedPeriods[number]; const status = tracker.generateBudgetStatus(period); if (options.json) { @@ -3204,7 +3219,9 @@ program const remaining = status.currentPeriod.remaining ?? 0; const remainingColor = remaining > status.currentPeriod.limit * 0.2 ? chalk.green : chalk.yellow; console.log(` Remaining: ${remainingColor(`$${remaining.toFixed(2)}`)}`); - const percentUsed = (status.currentPeriod.spent / status.currentPeriod.limit) * 100; + const percentUsed = status.currentPeriod.limit > 0 + ? (status.currentPeriod.spent / status.currentPeriod.limit) * 100 + : 0; console.log(` Usage: ${percentUsed.toFixed(1)}%`); } else { console.log(chalk.dim(' Limit: (not set)')); diff --git a/src/config.ts b/src/config.ts index 39ec2a0..9717245 100644 --- a/src/config.ts +++ b/src/config.ts @@ -98,11 +98,12 @@ export class ConfigManager { const parsed = JSON.parse(content); // Merge with defaults to ensure all required fields exist + // Use structuredClone to prevent shared references (e.g., warningThresholds array) return { ...structuredClone(DEFAULT_CONFIG), ...parsed, budget: { - ...DEFAULT_BUDGET_CONFIG, + ...structuredClone(DEFAULT_BUDGET_CONFIG), ...parsed.budget } }; @@ -147,10 +148,13 @@ export class ConfigManager { } if (this.pendingSave) { - // Only write if directory still exists (handles test cleanup race conditions) - const dir = dirname(this.configPath); - if (existsSync(dir)) { + try { writeFileSync(this.configPath, JSON.stringify(this.config, null, 2)); + } catch (err) { + // Ignore ENOENT if directory was removed (e.g., test cleanup race conditions) + if ((err as NodeJS.ErrnoException).code !== 'ENOENT') { + throw err; + } } this.pendingSave = false; } @@ -323,6 +327,13 @@ export class ConfigManager { } } + // Validate e2bHourlyRate + if (config.e2bHourlyRate !== undefined && config.e2bHourlyRate !== null) { + if (config.e2bHourlyRate < 0) { + throw new Error('E2B hourly rate must be a non-negative number'); + } + } + // Merge config this.config.budget = { ...this.config.budget, From f18f53fcfc5294d40b0b5b101c12c16e324501e3 Mon Sep 17 00:00:00 2001 From: Test User Date: Thu, 29 Jan 2026 22:50:11 -0700 Subject: [PATCH 5/7] fix: route budget keys through setBudgetConfig() for validation Budget keys set via 'config set budget.*' were bypassing validation in setBudgetConfig(). Now routes all budget.* keys through the validated method to ensure non-negative limits and valid thresholds. --- src/cli.ts | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/src/cli.ts b/src/cli.ts index e4cb7cd..1201c9a 100644 --- a/src/cli.ts +++ b/src/cli.ts @@ -42,7 +42,7 @@ import { existsSync } from 'fs'; import * as path from 'path'; import * as os from 'os'; import { randomUUID } from 'crypto'; -import { SandboxStatus, type E2BSession, type StatusResult, type SessionInfo } from './types.js'; +import { SandboxStatus, type BudgetConfig, type E2BSession, type StatusResult, type SessionInfo } from './types.js'; program .name('parallel-cc') @@ -3069,7 +3069,14 @@ configCmd parsedValue = value; } - configManager.set(normalizedKey, parsedValue); + // Route budget keys through setBudgetConfig() for validation + if (normalizedKey.startsWith('budget.')) { + const budgetKey = normalizedKey.replace(/^budget\./, ''); + const budgetConfig = { [budgetKey]: parsedValue } as Partial; + configManager.setBudgetConfig(budgetConfig); + } else { + configManager.set(normalizedKey, parsedValue); + } if (options.json) { console.log(JSON.stringify({ success: true, key: normalizedKey, value: parsedValue })); From 55cd91855170db58d0389627a6c781a392af0efe Mon Sep 17 00:00:00 2001 From: Test User Date: Thu, 29 Jan 2026 22:50:52 -0700 Subject: [PATCH 6/7] fix: backup corrupted config file before overwriting with defaults When JSON parsing fails, create a timestamped backup of the corrupted config file before returning defaults. This prevents silent data loss when the user's config has a syntax error. --- src/config.ts | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/src/config.ts b/src/config.ts index 9717245..9a48470 100644 --- a/src/config.ts +++ b/src/config.ts @@ -4,7 +4,7 @@ * Handles persistent user settings stored in ~/.parallel-cc/config.json */ -import { existsSync, mkdirSync, readFileSync, writeFileSync } from 'fs'; +import { copyFileSync, existsSync, mkdirSync, readFileSync, writeFileSync } from 'fs'; import { dirname, join as pathJoin } from 'path'; import { homedir } from 'os'; import type { BudgetConfig } from './types.js'; @@ -107,8 +107,16 @@ export class ConfigManager { ...parsed.budget } }; - } catch { - // Invalid JSON - return defaults + } catch (err) { + // Invalid JSON - backup corrupted file and return defaults + const backupPath = `${this.configPath}.corrupted.${Date.now()}`; + try { + copyFileSync(this.configPath, backupPath); + console.error(`Warning: Config file had invalid JSON. Backed up to: ${backupPath}`); + } catch { + // Backup failed - just warn + console.error(`Warning: Config file has invalid JSON and could not be backed up: ${(err as Error).message}`); + } return structuredClone(DEFAULT_CONFIG); } } From c3a8ddf94a81d47b721e860069a5f9d5660fdeab Mon Sep 17 00:00:00 2001 From: Test User Date: Thu, 29 Jan 2026 23:02:01 -0700 Subject: [PATCH 7/7] refactor: use 'budget status' subcommand per naming convention Per project guidelines, prefer proper subcommands over hyphenated names. - Created 'budget' command group with 'status' subcommand - Kept 'budget-status' as alias for backward compatibility - Updated CLAUDE.md CLI Commands table with budget and config commands --- CLAUDE.md | 5 +++++ src/cli.ts | 39 ++++++++++++++++++++++++++++----------- 2 files changed, 33 insertions(+), 11 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 756b6b9..c8f9d43 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -381,6 +381,11 @@ CREATE INDEX idx_budget_period ON budget_tracking(period, period_start); | `templates delete ` | Delete a custom template (v1.1) | | `templates export ` | Export a template to JSON (v1.1) | | `templates import` | Import a template from JSON (v1.1) | +| `config set ` | Set a configuration value (v1.1) | +| `config get ` | Get a configuration value (v1.1) | +| `config list` | Display all configuration values (v1.1) | +| `budget status` | Show budget and spending status (v1.1) | +| `budget-status` | Alias for `budget status` (backward compat) | ## MCP Server Tools (v0.4) diff --git a/src/cli.ts b/src/cli.ts index 1201c9a..646f3e2 100644 --- a/src/cli.ts +++ b/src/cli.ts @@ -3172,18 +3172,13 @@ configCmd }); // ============================================================================ -// Budget Status Command (v1.1) +// Budget Commands (v1.1) // ============================================================================ /** - * Budget status command - Show spending and budget information + * Shared action handler for budget status command */ -program - .command('budget-status') - .description('Show budget and spending status (v1.1)') - .option('--period ', 'Show specific period (daily, weekly, monthly)', 'monthly') - .option('--json', 'Output as JSON') - .action(async (options) => { +async function budgetStatusAction(options: { period: string; json?: boolean }) { let db: SessionDB | null = null; try { db = new SessionDB(); @@ -3198,7 +3193,7 @@ program const tracker = new BudgetTracker(db, configManager); // Validate period option - const allowedPeriods = ['daily', 'weekly', 'monthly'] as const; + const allowedPeriods: readonly string[] = ['daily', 'weekly', 'monthly']; if (!allowedPeriods.includes(options.period)) { const message = `Invalid period: ${options.period}. Use daily, weekly, or monthly.`; if (options.json) { @@ -3208,7 +3203,7 @@ program } process.exit(1); } - const period = options.period as typeof allowedPeriods[number]; + const period = options.period as 'daily' | 'weekly' | 'monthly'; const status = tracker.generateBudgetStatus(period); if (options.json) { @@ -3267,7 +3262,29 @@ program db.close(); } } - }); +} + +/** + * Budget command group - Manage budget settings and view status + */ +const budgetCmd = program + .command('budget') + .description('Manage budget settings and view status (v1.1)'); + +budgetCmd + .command('status') + .description('Show budget and spending status') + .option('--period ', 'Show specific period (daily, weekly, monthly)', 'monthly') + .option('--json', 'Output as JSON') + .action(budgetStatusAction); + +// Backward compatibility alias: budget-status -> budget status +program + .command('budget-status') + .description('Show budget and spending status (alias for "budget status")') + .option('--period ', 'Show specific period (daily, weekly, monthly)', 'monthly') + .option('--json', 'Output as JSON') + .action(budgetStatusAction); // Parse and execute program.parse();