diff --git a/apps/cli/src/commands/eval/run-eval.ts b/apps/cli/src/commands/eval/run-eval.ts index 939f33b2a..71d0aae68 100644 --- a/apps/cli/src/commands/eval/run-eval.ts +++ b/apps/cli/src/commands/eval/run-eval.ts @@ -293,6 +293,7 @@ async function prepareFileMetadata(params: { readonly suiteTargets?: readonly string[]; readonly yamlCache?: boolean; readonly yamlCachePath?: string; + readonly totalBudgetUsd?: number; }> { const { testFilePath, repoRoot, cwd, options } = params; @@ -383,6 +384,7 @@ async function prepareFileMetadata(params: { suiteTargets, yamlCache: suite.cacheConfig?.enabled, yamlCachePath: suite.cacheConfig?.cachePath, + totalBudgetUsd: suite.totalBudgetUsd, }; } @@ -423,6 +425,7 @@ async function runSingleEvalFile(params: { readonly evalCases: readonly EvalTest[]; readonly trialsConfig?: TrialsConfig; readonly matrixMode?: boolean; + readonly totalBudgetUsd?: number; }): Promise<{ results: EvaluationResult[] }> { const { testFilePath, @@ -442,6 +445,7 @@ async function runSingleEvalFile(params: { evalCases, trialsConfig, matrixMode, + totalBudgetUsd, } = params; const targetName = selection.targetName; @@ -523,6 +527,7 @@ async function runSingleEvalFile(params: { keepWorkspaces: options.keepWorkspaces, cleanupWorkspaces: options.cleanupWorkspaces, trials: trialsConfig, + totalBudgetUsd, streamCallbacks: streamingObserver?.getStreamCallbacks(), onResult: async (result: EvaluationResult) => { // Finalize streaming observer span with score @@ -721,6 +726,7 @@ export async function runEvalCommand(input: RunEvalCommandInput): Promise readonly suiteTargets?: readonly string[]; readonly yamlCache?: boolean; readonly yamlCachePath?: string; + readonly totalBudgetUsd?: number; } >(); for (const testFilePath of resolvedTestFiles) { @@ -861,6 +867,7 @@ export async function runEvalCommand(input: RunEvalCommandInput): Promise evalCases: applicableEvalCases, trialsConfig: targetPrep.trialsConfig, matrixMode: targetPrep.selections.length > 1, + totalBudgetUsd: targetPrep.totalBudgetUsd, }); allResults.push(...result.results); diff --git a/packages/core/src/evaluation/loaders/config-loader.ts b/packages/core/src/evaluation/loaders/config-loader.ts index a1970ef98..1341ff1a4 100644 --- a/packages/core/src/evaluation/loaders/config-loader.ts +++ b/packages/core/src/evaluation/loaders/config-loader.ts @@ -262,6 +262,33 @@ export function extractCacheConfig(suite: JsonObject): CacheConfig | undefined { return { enabled: cache, cachePath: resolvedCachePath }; } +/** + * Extract suite-level total budget from parsed eval suite's execution block. + * Returns undefined when not specified. + */ +export function extractTotalBudgetUsd(suite: JsonObject): number | undefined { + const execution = suite.execution; + if (!execution || typeof execution !== 'object' || Array.isArray(execution)) { + return undefined; + } + + const executionObj = execution as Record; + const rawBudget = executionObj.total_budget_usd ?? executionObj.totalBudgetUsd; + + if (rawBudget === undefined || rawBudget === null) { + return undefined; + } + + if (typeof rawBudget === 'number' && rawBudget > 0) { + return rawBudget; + } + + logWarning( + `Invalid execution.total_budget_usd: ${rawBudget}. Must be a positive number. Ignoring.`, + ); + return undefined; +} + function logWarning(message: string): void { console.warn(`${ANSI_YELLOW}Warning: ${message}${ANSI_RESET}`); } diff --git a/packages/core/src/evaluation/orchestrator.ts b/packages/core/src/evaluation/orchestrator.ts index e915ba1b3..61fb04ae1 100644 --- a/packages/core/src/evaluation/orchestrator.ts +++ b/packages/core/src/evaluation/orchestrator.ts @@ -153,6 +153,8 @@ export interface RunEvaluationOptions { readonly trials?: TrialsConfig; /** Real-time observability callbacks passed to the provider */ readonly streamCallbacks?: ProviderStreamCallbacks; + /** Suite-level total cost budget in USD (stops dispatching when exceeded) */ + readonly totalBudgetUsd?: number; } export async function runEvaluation( @@ -179,6 +181,7 @@ export async function runEvaluation( cleanupWorkspaces, trials, streamCallbacks, + totalBudgetUsd, } = options; // Disable cache when trials > 1 (cache makes trials deterministic = pointless) @@ -403,6 +406,10 @@ export async function runEvaluation( const workerIdByEvalId = new Map(); let beforeAllOutputAttached = false; + // Suite-level budget tracking + let cumulativeBudgetCost = 0; + let budgetExhausted = false; + // Map test cases to limited promises for parallel execution const promises = filteredEvalCases.map((evalCase) => limit(async () => { @@ -410,6 +417,36 @@ export async function runEvaluation( const workerId = nextWorkerId++; workerIdByEvalId.set(evalCase.id, workerId); + // Check suite-level budget before dispatching + if (totalBudgetUsd !== undefined && budgetExhausted) { + const budgetResult: EvaluationResult = { + timestamp: (now ?? (() => new Date()))().toISOString(), + testId: evalCase.id, + dataset: evalCase.dataset, + score: 0, + hits: [], + misses: [], + answer: '', + target: target.name, + error: `Suite budget exceeded ($${cumulativeBudgetCost.toFixed(4)} / $${totalBudgetUsd.toFixed(4)})`, + budgetExceeded: true, + }; + + if (onProgress) { + await onProgress({ + workerId, + testId: evalCase.id, + status: 'failed', + completedAt: Date.now(), + error: budgetResult.error, + }); + } + if (onResult) { + await onResult(budgetResult); + } + return budgetResult; + } + if (onProgress) { await onProgress({ workerId, @@ -448,6 +485,26 @@ export async function runEvaluation( ? await runEvalCaseWithTrials(runCaseOptions, trials) : await runEvalCase(runCaseOptions); + // Track suite-level budget + if (totalBudgetUsd !== undefined) { + // Sum all trial costs when trials are used, otherwise use trace cost + let caseCost: number | undefined; + if (result.trials && result.trials.length > 0) { + const trialCostSum = result.trials.reduce((sum, t) => sum + (t.costUsd ?? 0), 0); + if (trialCostSum > 0) { + caseCost = trialCostSum; + } + } else { + caseCost = result.trace?.costUsd; + } + if (caseCost !== undefined) { + cumulativeBudgetCost += caseCost; + if (cumulativeBudgetCost >= totalBudgetUsd) { + budgetExhausted = true; + } + } + } + // Attach beforeAllOutput to first result only if (beforeAllOutput && !beforeAllOutputAttached) { result = { ...result, beforeAllOutput }; diff --git a/packages/core/src/evaluation/types.ts b/packages/core/src/evaluation/types.ts index 87f296cbc..234734036 100644 --- a/packages/core/src/evaluation/types.ts +++ b/packages/core/src/evaluation/types.ts @@ -716,6 +716,8 @@ export interface EvaluationResult { readonly aggregation?: TrialAggregation; /** Whether the trial loop was terminated early due to cost limit */ readonly costLimited?: boolean; + /** Whether the evaluation was skipped due to suite-level budget exhaustion */ + readonly budgetExceeded?: boolean; } export type EvaluationVerdict = 'pass' | 'fail' | 'borderline'; diff --git a/packages/core/src/evaluation/yaml-parser.ts b/packages/core/src/evaluation/yaml-parser.ts index b903234a2..e3f8e704c 100644 --- a/packages/core/src/evaluation/yaml-parser.ts +++ b/packages/core/src/evaluation/yaml-parser.ts @@ -9,6 +9,7 @@ import { extractTargetFromSuite, extractTargetsFromSuite, extractTargetsFromTestCase, + extractTotalBudgetUsd, extractTrialsConfig, loadConfig, } from './loaders/config-loader.js'; @@ -153,6 +154,8 @@ export type EvalSuiteResult = { readonly cacheConfig?: import('./loaders/config-loader.js').CacheConfig; /** Suite-level metadata (name, description, version, etc.) */ readonly metadata?: import('./metadata.js').EvalMetadata; + /** Suite-level total cost budget in USD */ + readonly totalBudgetUsd?: number; }; /** @@ -175,6 +178,7 @@ export async function loadTestSuite( trials: extractTrialsConfig(parsed), targets: extractTargetsFromSuite(parsed), cacheConfig: extractCacheConfig(parsed), + totalBudgetUsd: extractTotalBudgetUsd(parsed), ...(metadata !== undefined && { metadata }), }; } diff --git a/packages/core/test/evaluation/loaders/config-loader.test.ts b/packages/core/test/evaluation/loaders/config-loader.test.ts index 292b3aae9..3d551d504 100644 --- a/packages/core/test/evaluation/loaders/config-loader.test.ts +++ b/packages/core/test/evaluation/loaders/config-loader.test.ts @@ -4,6 +4,7 @@ import { extractTargetFromSuite, extractTargetsFromSuite, extractTargetsFromTestCase, + extractTotalBudgetUsd, extractTrialsConfig, } from '../../../src/evaluation/loaders/config-loader.js'; import type { JsonObject } from '../../../src/evaluation/types.js'; @@ -224,3 +225,40 @@ describe('extractTargetsFromTestCase', () => { expect(extractTargetsFromTestCase(testCase)).toBeUndefined(); }); }); + +describe('extractTotalBudgetUsd', () => { + it('returns undefined when no execution block', () => { + const suite: JsonObject = { tests: [] }; + expect(extractTotalBudgetUsd(suite)).toBeUndefined(); + }); + + it('returns undefined when no total_budget_usd in execution', () => { + const suite: JsonObject = { execution: { target: 'default' } }; + expect(extractTotalBudgetUsd(suite)).toBeUndefined(); + }); + + it('parses valid total_budget_usd (snake_case)', () => { + const suite: JsonObject = { execution: { total_budget_usd: 10.0 } }; + expect(extractTotalBudgetUsd(suite)).toBe(10.0); + }); + + it('parses valid totalBudgetUsd (camelCase)', () => { + const suite: JsonObject = { execution: { totalBudgetUsd: 5.5 } }; + expect(extractTotalBudgetUsd(suite)).toBe(5.5); + }); + + it('returns undefined for zero budget', () => { + const suite: JsonObject = { execution: { total_budget_usd: 0 } }; + expect(extractTotalBudgetUsd(suite)).toBeUndefined(); + }); + + it('returns undefined for negative budget', () => { + const suite: JsonObject = { execution: { total_budget_usd: -1 } }; + expect(extractTotalBudgetUsd(suite)).toBeUndefined(); + }); + + it('returns undefined for non-number budget', () => { + const suite: JsonObject = { execution: { total_budget_usd: 'ten' } }; + expect(extractTotalBudgetUsd(suite)).toBeUndefined(); + }); +}); diff --git a/packages/core/test/evaluation/orchestrator.test.ts b/packages/core/test/evaluation/orchestrator.test.ts index fb977d778..99dbd83aa 100644 --- a/packages/core/test/evaluation/orchestrator.test.ts +++ b/packages/core/test/evaluation/orchestrator.test.ts @@ -2049,3 +2049,160 @@ describe('workspace.template .code-workspace resolution', () => { } }); }); + +describe('suite-level total budget guardrail', () => { + it('completes normally when totalBudgetUsd is not set', async () => { + const provider: Provider = { + id: 'budget:mock', + kind: 'mock' as const, + targetName: 'mock', + async invoke(): Promise { + return { + output: [{ role: 'assistant', content: 'response' }], + costUsd: 1.0, + }; + }, + }; + + const evalCases: EvalTest[] = [ + { ...baseTestCase, id: 'case-1' }, + { ...baseTestCase, id: 'case-2' }, + { ...baseTestCase, id: 'case-3' }, + ]; + + const results = await runEvaluation({ + testFilePath: 'in-memory.yaml', + repoRoot: 'in-memory', + target: baseTarget, + providerFactory: () => provider, + evaluators: evaluatorRegistry, + evalCases, + }); + + expect(results).toHaveLength(3); + expect(results.every((r) => r.budgetExceeded === undefined)).toBe(true); + }); + + it('completes normally when budget is not exceeded', async () => { + const provider: Provider = { + id: 'budget:mock', + kind: 'mock' as const, + targetName: 'mock', + async invoke(): Promise { + return { + output: [{ role: 'assistant', content: 'response' }], + costUsd: 1.0, + }; + }, + }; + + const evalCases: EvalTest[] = [ + { ...baseTestCase, id: 'case-1' }, + { ...baseTestCase, id: 'case-2' }, + ]; + + const results = await runEvaluation({ + testFilePath: 'in-memory.yaml', + repoRoot: 'in-memory', + target: baseTarget, + providerFactory: () => provider, + evaluators: evaluatorRegistry, + evalCases, + totalBudgetUsd: 10.0, + }); + + expect(results).toHaveLength(2); + expect(results.every((r) => r.budgetExceeded === undefined)).toBe(true); + }); + + it('stops dispatching when budget is exceeded mid-run', async () => { + const provider: Provider = { + id: 'budget:mock', + kind: 'mock' as const, + targetName: 'mock', + async invoke(): Promise { + return { + output: [{ role: 'assistant', content: 'response' }], + costUsd: 3.0, + }; + }, + }; + + const evalCases: EvalTest[] = [ + { ...baseTestCase, id: 'case-1' }, + { ...baseTestCase, id: 'case-2' }, + { ...baseTestCase, id: 'case-3' }, + { ...baseTestCase, id: 'case-4' }, + ]; + + const results = await runEvaluation({ + testFilePath: 'in-memory.yaml', + repoRoot: 'in-memory', + target: baseTarget, + providerFactory: () => provider, + evaluators: evaluatorRegistry, + evalCases, + totalBudgetUsd: 5.0, + maxConcurrency: 1, + }); + + expect(results).toHaveLength(4); + + // First two should run normally ($3 + $3 = $6 >= $5) + expect(results[0].budgetExceeded).toBeUndefined(); + expect(results[1].budgetExceeded).toBeUndefined(); + + // Remaining should be budget-exceeded + expect(results[2].budgetExceeded).toBe(true); + expect(results[3].budgetExceeded).toBe(true); + expect(results[2].error).toContain('Suite budget exceeded'); + expect(results[3].error).toContain('Suite budget exceeded'); + expect(results[2].score).toBe(0); + expect(results[3].score).toBe(0); + }); + + it('works correctly with trials and budget', async () => { + const provider: Provider = { + id: 'budget:mock', + kind: 'mock' as const, + targetName: 'mock', + async invoke(): Promise { + return { + output: [{ role: 'assistant', content: 'response' }], + costUsd: 2.0, + }; + }, + }; + + const evalCases: EvalTest[] = [ + { ...baseTestCase, id: 'case-1' }, + { ...baseTestCase, id: 'case-2' }, + { ...baseTestCase, id: 'case-3' }, + { ...baseTestCase, id: 'case-4' }, + ]; + + // evaluatorRegistry always returns 0.8 (pass), so pass_at_k exits after 1 trial per case. + // Each case costs $2. Budget of $5 is exceeded after case-3 ($6 >= $5). + // Case-4 should be budget-exceeded. + const results = await runEvaluation({ + testFilePath: 'in-memory.yaml', + repoRoot: 'in-memory', + target: baseTarget, + providerFactory: () => provider, + evaluators: evaluatorRegistry, + evalCases, + totalBudgetUsd: 5.0, + maxConcurrency: 1, + trials: { count: 2, strategy: 'pass_at_k' }, + }); + + expect(results).toHaveLength(4); + // First three run normally + expect(results[0].budgetExceeded).toBeUndefined(); + expect(results[1].budgetExceeded).toBeUndefined(); + expect(results[2].budgetExceeded).toBeUndefined(); + // Fourth should be budget-exceeded + expect(results[3].budgetExceeded).toBe(true); + expect(results[3].error).toContain('Suite budget exceeded'); + }); +});