Repository Public
1import { afterAll, beforeAll, describe, expect, it } from 'vitest';2import { BudgetExhaustedError, type RunContext, type RunLogic } from '../src/index.js';3import { createTestDb, newMission, rows, startWorker, steps, waitStatus, type TestDb } from './helpers.js';45let db: TestDb;6beforeAll(async () => {7 db = await createTestDb();8});9afterAll(async () => {10 await db.cleanup();11});1213const env = (status: 'done' | 'budget_exhausted', payload: unknown = null) => ({ status, summary: status, payload, artifacts: [], evidence: [], openQuestions: [] });14const call = (ctx: RunContext, n: number, usd: number, counter?: { n: number }) =>15 ctx.step({ kind: 'model_call', name: 'model', effect: 'model_call', request: { n } }, async () => {16 if (counter) counter.n++;17 return { result: { n }, usage: { usd, inputTokens: 1000, outputTokens: 100, cacheReadTokens: 10, cacheWriteTokens: 5 } };18 });19const spendToday = async (): Promise<number> => Number((await rows<{ usd: string }>(db, 'SELECT usd FROM spend_days WHERE day = current_date'))[0]?.usd ?? 0);2021describe('budgets and usage accounting', () => {22 it('accounts usage in runs.usage, missions.spent_usd and spend_days, and reports it through ctx', async () => {23 const before = await spendToday();24 const { missionId, rootRunId } = await newMission(db, { budgetUsd: 1, limits: { turns: 10, wallSec: 100 } });25 const w = startWorker(db, async (ctx) => {26 const l0 = ctx.budgetLeft();27 await call(ctx, 1, 0.1);28 await call(ctx, 2, 0.15);29 await ctx.step({ kind: 'tool_call', name: 't', effect: 'read', request: {} }, async () => ({ result: {} }));30 return env('done', { l0, usage: ctx.usage(), left: ctx.budgetLeft() });31 });32 const run = await waitStatus(db, rootRunId, ['done']);33 await w.stop();34 const p = run.result?.payload as { l0: { usd: number; turns: number; wallSec: number }; usage: Record<string, number>; left: { usd: number; turns: number; wallSec: number } };35 expect(p.l0).toEqual({ usd: 1, turns: 10, wallSec: 100 });36 expect(p.usage).toEqual({ usd: 0.25, inputTokens: 2000, outputTokens: 200, cacheReadTokens: 20, cacheWriteTokens: 10, turns: 2, toolCalls: 1 });37 expect(p.left.usd).toBeCloseTo(0.75);38 expect(p.left.turns).toBe(8);39 expect(p.left.wallSec).toBeLessThanOrEqual(100);40 expect(p.left.wallSec).toBeGreaterThan(90);41 expect(run.usage).toEqual(p.usage);42 expect(Number((await rows<{ spent_usd: string }>(db, 'SELECT spent_usd FROM missions WHERE id = $1', [missionId]))[0]?.spent_usd)).toBeCloseTo(0.25);43 expect((await spendToday()) - before).toBeCloseTo(0.25);44 const usageEvents = await rows<{ payload: { usage: { usd: number }; missionSpentUsd: number } }>(db, `SELECT payload FROM events WHERE run_id = $1 AND kind = 'run.usage' ORDER BY id`, [rootRunId]);45 expect(usageEvents.map((e) => e.payload.usage.usd)).toEqual([0.1, 0.25, 0.25]);46 });4748 it('hard stop: a model_call is refused once usd ≤ 0; uncaught ⇒ run budget_exhausted', async () => {49 const { missionId, rootRunId } = await newMission(db, { budgetUsd: 0.05 });50 const counter = { n: 0 };51 const w = startWorker(db, async (ctx) => {52 for (let n = 0; n < 10; n++) await call(ctx, n, 0.03, counter);53 return env('done');54 });55 const run = await waitStatus(db, rootRunId, ['budget_exhausted', 'done', 'failed']);56 await w.stop();57 expect(run.status).toBe('budget_exhausted');58 expect(counter.n).toBe(2); // 0.03 + 0.03 ≥ 0.05 ⇒ the third call never executes59 expect(run.usage.usd).toBeCloseTo(0.06);60 expect(run.result?.summary).toMatch(/run_budget/);61 const st = await steps(db, rootRunId);62 expect(st.map((s) => s.status)).toEqual(['completed', 'completed', 'failed']);63 expect(st[2]?.error?.name).toBe('BudgetExhaustedError');64 expect((await rows<{ status: string }>(db, 'SELECT status FROM missions WHERE id = $1', [missionId]))[0]?.status).toBe('budget_exhausted');65 });6667 it('the refusal is journaled: logic that catches it replays identically, non-model steps still run', async () => {68 const { rootRunId } = await newMission(db, { budgetUsd: 0.05 });69 const counter = { n: 0 };70 const seen: string[] = [];71 const logic: RunLogic = async (ctx) => {72 let stopped: string | null = null;73 let turns = 0;74 while (!stopped) {75 if (ctx.budgetLeft().usd <= 0) {76 // The engine would stop here; call anyway to prove the kernel refuses.77 try {78 await call(ctx, turns, 0.03, counter);79 } catch (err) {80 stopped = err instanceof BudgetExhaustedError ? `${err.name}:${err.reason}` : String(err);81 }82 } else {83 await call(ctx, turns++, 0.03, counter);84 }85 }86 seen.push(`${stopped}|${turns}|${ctx.budgetLeft().usd.toFixed(2)}`);87 await ctx.step({ kind: 'tool_call', name: 'git_commit', effect: 'idempotent_write', request: {} }, async () => ({ result: { ok: true } }));88 await ctx.sleep(300); // suspend ⇒ replay through the refusal89 seen.push('after-sleep');90 return env('budget_exhausted', { stopped, turns });91 };92 const w = startWorker(db, logic);93 const run = await waitStatus(db, rootRunId, ['budget_exhausted', 'failed']);94 await w.stop();95 expect(run.status).toBe('budget_exhausted');96 expect(run.result?.payload).toEqual({ stopped: 'BudgetExhaustedError:run_budget', turns: 2 });97 expect(seen).toEqual(['BudgetExhaustedError:run_budget|2|-0.01', 'BudgetExhaustedError:run_budget|2|-0.01', 'after-sleep']);98 expect(counter.n).toBe(2);99 });100101 it('daily cap: once today’s spend reaches the cap every run’s model_call is refused', async () => {102 const cap = (await spendToday()) + 0.05;103 const a = await newMission(db, { budgetUsd: 5 });104 const counter = { n: 0 };105 const logic: RunLogic = async (ctx) => {106 for (let n = 0; n < 5; n++) await call(ctx, n, 0.03, counter);107 return env('done');108 };109 const w = startWorker(db, logic, { dailyBudgetUsd: () => cap });110 const ra = await waitStatus(db, a.rootRunId, ['budget_exhausted', 'done', 'failed']);111 expect(ra.status).toBe('budget_exhausted');112 expect(ra.result?.summary).toMatch(/daily_cap/);113 expect(counter.n).toBe(2);114 const b = await newMission(db, { budgetUsd: 5 }); // a fresh mission with plenty of budget115 const rb = await waitStatus(db, b.rootRunId, ['budget_exhausted', 'done', 'failed']);116 await w.stop();117 expect(rb.status).toBe('budget_exhausted');118 expect(counter.n).toBe(2); // nothing executed for the second mission119 expect(rb.usage.usd).toBe(0);120 });121122 it('default daily cap comes from the owner’s settings (application_settings.dailyBudgetUsd)', async () => {123 const today = await spendToday();124 await db.k.pool.query(125 `INSERT INTO application_settings (user_id, settings) VALUES ($1, $2::jsonb) ON CONFLICT (user_id) DO UPDATE SET settings = EXCLUDED.settings`,126 [db.ownerId, JSON.stringify({ dailyBudgetUsd: Math.max(today, 0.000001) })],127 );128 const { rootRunId } = await newMission(db);129 const counter = { n: 0 };130 const w = startWorker(db, async (ctx) => {131 await call(ctx, 0, 0.01, counter);132 return env('done');133 }, { dailyBudgetUsd: undefined });134 const run = await waitStatus(db, rootRunId, ['budget_exhausted', 'done', 'failed']);135 await w.stop();136 await db.k.pool.query('DELETE FROM application_settings WHERE user_id = $1', [db.ownerId]);137 expect(run.status).toBe('budget_exhausted');138 expect(counter.n).toBe(0);139 });140});141