| 1 | import assert from "node:assert/strict"; |
| 2 | import { test } from "node:test"; |
| 3 | |
| 4 | import { agentStatus, budgetBlock, chargedMicros, costMicros, dayKey, monthKey, replyCapMicros, totalTokens } from "./budget.ts"; |
| 5 | |
| 6 | const OCT_8 = new Date("2026-10-08T15:00:00Z"); |
| 7 | const none = { monthly_micros: null, daily_micros: null, task_micros: null }; |
| 8 | |
| 9 | test("spend rolls up by UTC month and day", () => { |
| 10 | assert.equal(monthKey(OCT_8), "2026-10"); |
| 11 | assert.equal(dayKey(OCT_8), "2026-10-08"); |
| 12 | assert.equal(dayKey(new Date("2026-10-31T23:59:59Z")), "2026-10-31"); |
| 13 | }); |
| 14 | |
| 15 | test("no caps, no block; a cap of zero is no cap", () => { |
| 16 | assert.equal(budgetBlock(none, { month: 9e9, day: 9e9 }, OCT_8), null); |
| 17 | assert.equal(budgetBlock({ ...none, monthly_micros: 0 }, { month: 9e9, day: 0 }, OCT_8), null); |
| 18 | }); |
| 19 | |
| 20 | test("the monthly cap blocks at 100%, and says so in plain words", () => { |
| 21 | const budget = { ...none, monthly_micros: 5_000_000 }; |
| 22 | assert.equal(budgetBlock(budget, { month: 4_999_999, day: 0 }, OCT_8), null); |
| 23 | const blocked = budgetBlock(budget, { month: 5_000_000, day: 0 }, OCT_8); |
| 24 | assert.equal(blocked?.cap, "month"); |
| 25 | assert.equal(blocked?.message, "I'm out of budget for October. An owner can raise my monthly limit on my profile."); |
| 26 | }); |
| 27 | |
| 28 | test("the daily cap blocks too, after the monthly one is checked", () => { |
| 29 | const budget = { ...none, monthly_micros: 5_000_000, daily_micros: 1_000_000 }; |
| 30 | assert.equal(budgetBlock(budget, { month: 2_000_000, day: 1_000_000 }, OCT_8)?.cap, "day"); |
| 31 | assert.equal(budgetBlock(budget, { month: 5_000_000, day: 1_000_000 }, OCT_8)?.cap, "month"); |
| 32 | }); |
| 33 | |
| 34 | test("a reply's cap is the lowest of what is left and the per-task and plan caps", () => { |
| 35 | assert.equal(replyCapMicros(none, { month: 0, day: 0 }, null), null); |
| 36 | assert.equal(replyCapMicros({ ...none, monthly_micros: 1_000_000 }, { month: 900_000, day: 0 }, null), 100_000); |
| 37 | assert.equal(replyCapMicros({ ...none, monthly_micros: 1_000_000, task_micros: 50_000 }, { month: 0, day: 0 }, 2_000_000), 50_000); |
| 38 | assert.equal(replyCapMicros(none, { month: 0, day: 0 }, 30_000), 30_000, "the plan's per-run cap"); |
| 39 | assert.equal(replyCapMicros({ ...none, daily_micros: 10 }, { month: 0, day: 50 }, null), 1, "never read as no cap"); |
| 40 | }); |
| 41 | |
| 42 | test("cost is tokens at the provider's price per million, rounded up to a micro-dollar", () => { |
| 43 | const price = { input: 0.1, output: 0.5, cacheRead: 0.01, cacheWrite: 0.125 }; |
| 44 | const tokens = { input: 2_000, output: 300, cacheRead: 1_000, cacheWrite: 0 }; |
| 45 | // 2000*0.1 + 300*0.5 + 1000*0.01 = 200 + 150 + 10 |
| 46 | assert.equal(costMicros(tokens, price), 360); |
| 47 | assert.equal(costMicros(tokens, null), 0, "an unpriced model costs g1t nothing here"); |
| 48 | assert.equal(totalTokens(tokens), 3_300); |
| 49 | }); |
| 50 | |
| 51 | test("what a reply counts: the model with its margin on g1t's models, the agent rate always", () => { |
| 52 | // $0.25 per million tokens is 250,000 micros per million. |
| 53 | assert.equal(chargedMicros({ costMicros: 1_000, hosted: true, marginPercent: 0, ratePerMillionMicros: 250_000, tokens: 4_000 }), 2_000); |
| 54 | assert.equal(chargedMicros({ costMicros: 1_000, hosted: true, marginPercent: 20, ratePerMillionMicros: 0, tokens: 4_000 }), 1_200); |
| 55 | assert.equal(chargedMicros({ costMicros: 1_000, hosted: false, marginPercent: 20, ratePerMillionMicros: 250_000, tokens: 4_000 }), 1_000, "own provider: the agent rate only"); |
| 56 | }); |
| 57 | |
| 58 | test("status: working while a reply is in flight, out of budget when capped, else idle", () => { |
| 59 | const later = new Date(OCT_8.getTime() + 60_000).toISOString(); |
| 60 | const earlier = new Date(OCT_8.getTime() - 60_000).toISOString(); |
| 61 | assert.equal(agentStatus({ busyUntil: later, now: OCT_8, blocked: true }), "working"); |
| 62 | assert.equal(agentStatus({ busyUntil: earlier, now: OCT_8, blocked: false }), "idle", "a desk that died is not working forever"); |
| 63 | assert.equal(agentStatus({ busyUntil: null, now: OCT_8, blocked: true }), "out_of_budget"); |
| 64 | }); |