Skip to content
Source
Sign inSign up
Open History to see the last commit

usageAggregation.test.ts

285 lines8.9 KB

1import { describe, expect, it } from "@effect/vitest";
2​
3import { resolveModelAliases, UsageAggregator } from "./usageAggregation.ts";
4import type { RateTable } from "./usagePricing.ts";
5import type { UsageRecord } from "@t3tools/provider-core/server/usage";
6​
7const rates: RateTable = new Map([
8 [
9 "claude-fable-5",
10 {
11 inputCostPerToken: 1e-5,
12 outputCostPerToken: 5e-5,
13 cacheReadCostPerToken: 1e-6,
14 cacheCreationCostPerToken: 1.25e-5,
15 cacheCreation1hCostPerToken: 2e-5,
16 fast: {
17 inputCostPerToken: 2e-5,
18 outputCostPerToken: 1e-4,
19 cacheReadCostPerToken: 2e-6,
20 cacheCreationCostPerToken: 2.5e-5,
21 cacheCreation1hCostPerToken: 4e-5,
22 },
23 ultrafast: null,
24 contextTiers: [],
25 },
26 ],
27]);
28​
29function record(overrides: Partial<UsageRecord> = {}): UsageRecord {
30 return {
31 provider: "claude",
32 timestampMs: Date.parse("2026-08-07T04:05:13.944Z"),
33 model: "claude-fable-5",
34 sessionId: "session-a",
35 totals: {
36 uncachedInputTokens: 100,
37 cachedInputTokens: 1000,
38 cacheCreationTokens: 10,
39 outputTokens: 50,
40 reasoningTokens: 0,
41 },
42 reportedCostUsd: null,
43 speed: "standard",
44 dedupeKey: null,
45 ...overrides,
46 };
48​
49function aggregate(
50 records: readonly UsageRecord[],
51 timeZone = "UTC",
52 resolution: "day" | "hour" = "day",
53) {
54 const hourlyBounds =
55 resolution === "hour"
56 ? {
57 sinceTimeMs: Date.parse("2026-08-06T04:37:00.000Z"),
58 untilTimeMs: Date.parse("2026-08-07T04:37:00.000Z"),
59 }
60 : {};
61 const aggregator = new UsageAggregator({
62 timeZone,
63 sinceDay: "2026-08-01",
64 untilDay: "2026-08-31",
65 resolution,
66 ...hourlyBounds,
67 rates,
68 });
69 for (const item of records) aggregator.add(item);
70 return aggregator.finish();
72​
73describe("UsageAggregator", () => {
74 it("splits a bucket's cost by category and speed", () => {
75 const [bucket] = aggregate([record(), record({ speed: "fast" })]).buckets;
76​
77 expect(bucket).toMatchObject({
78 costUsd: expect.closeTo(0.016875),
79 categoryCostUsd: {
80 input: expect.closeTo(0.003),
81 cacheRead: expect.closeTo(0.003),
82 cacheWrite: expect.closeTo(0.000375),
83 output: expect.closeTo(0.0075),
84 },
85 fastCostUsd: expect.closeTo(0.01125),
86 speedPremiumUsd: expect.closeTo(0.005625),
87 });
88 expect(bucket).not.toHaveProperty("ultrafastCostUsd");
89 });
90​
91 it("merges partial copies of a key by channel maximum in the first copy's bucket", () => {
92 const zeroTotals = {
93 uncachedInputTokens: 0,
94 cachedInputTokens: 0,
95 cacheCreationTokens: 0,
96 outputTokens: 0,
97 reasoningTokens: 0,
98 };
99 const partial = record({
100 dedupeKey: "msg_1:",
101 speed: "fast",
102 reportedCostUsd: 1,
103 totals: { ...record().totals, cachedInputTokens: 1200, outputTokens: 3, reasoningTokens: 5 },
104 });
105 const complete = record({
106 dedupeKey: "msg_1:",
107 model: "claude-fable-5[1m]",
108 timestampMs: Date.parse("2026-08-08T12:00:00Z"),
109 sessionId: "session-fork",
110 reportedCostUsd: 2,
111 totals: {
112 uncachedInputTokens: 2,
113 cachedInputTokens: 1000,
114 cacheCreationTokens: 40,
115 outputTokens: 264,
116 reasoningTokens: 135,
117 },
118 });
119 const zero = { ...partial, totals: zeroTotals, reportedCostUsd: null };
120 const merged = {
121 reportedCostUsd: 2,
122 totals: { ...complete.totals, uncachedInputTokens: 100, cachedInputTokens: 1200 },
123 };
124​
125 for (const copies of [
126 [zero, partial, complete, complete],
127 [complete, partial, zero, zero],
128 ]) {
129 const aggregator = new UsageAggregator({
130 timeZone: "UTC",
131 sinceDay: "2026-08-01",
132 untilDay: "2026-08-31",
133 rates,
134 });
135 copies.forEach((copy, index) => aggregator.add(copy, `${index}.jsonl`));
136 const result = aggregator.finish();
137 const expected = aggregate([{ ...copies[0]!, ...merged }]).buckets[0];
138​
139 expect(result.duplicatesDropped).toBe(3);
140 expect(result.buckets).toEqual([{ ...expected, sourcePath: "0.jsonl" }]);
141 }
142 });
143​
144 it("does not count a key whose first copy fell outside the window", () => {
145 const result = aggregate([
146 record({ dedupeKey: "msg_1:", timestampMs: Date.parse("2026-07-01T12:00:00Z") }),
147 record({ dedupeKey: "msg_1:", totals: { ...record().totals, outputTokens: 500 } }),
148 ]);
149​
150 expect(result.duplicatesDropped).toBe(1);
151 expect(result.outOfWindow).toBe(1);
152 expect(result.buckets).toHaveLength(0);
153 });
154​
155 it("buckets by the day in the requested time zone", () => {
156 const utc = aggregate([record()], "UTC");
157 const losAngeles = aggregate([record()], "America/Los_Angeles");
158​
159 expect(utc.buckets[0]?.day).toBe("2026-08-07");
160 expect(losAngeles.buckets[0]?.day).toBe("2026-08-06");
161 });
162​
163 it("finds the day at a quarter-hour zone's midnight across interleaved buckets", () => {
164 const result = aggregate(
165 [
166 record({ timestampMs: Date.parse("2026-08-01T18:14:59.999Z") }),
167 record({ timestampMs: Date.parse("2026-08-01T18:15:00.000Z") }),
168 record({ timestampMs: Date.parse("2026-08-01T18:15:00.000Z"), model: "claude-opus-5" }),
169 record({ timestampMs: Date.parse("2026-08-01T18:16:00.000Z") }),
170 ],
171 "Asia/Kathmandu",
172 );
173​
174 expect(result.buckets.map((bucket) => [bucket.day, bucket.model, bucket.records])).toEqual([
175 ["2026-08-01", "claude-fable-5", 1],
176 ["2026-08-02", "claude-fable-5", 2],
177 ["2026-08-02", "claude-opus-5", 1],
178 ]);
179 });
180​
181 it("finds the day at a fixed offset's midnight between quarter hours", () => {
182 const result = aggregate(
183 [
184 record({ timestampMs: Date.parse("2026-08-01T23:58:59.999Z") }),
185 record({ timestampMs: Date.parse("2026-08-01T23:59:00.000Z") }),
186 ],
187 "+00:01",
188 );
189​
190 expect(result.buckets.map((bucket) => [bucket.day, bucket.records])).toEqual([
191 ["2026-08-01", 1],
192 ["2026-08-02", 1],
193 ]);
194 });
195​
196 it("splits an hourly request into fixed buckets anchored to its exact start", () => {
197 const result = aggregate(
198 [
199 record({ timestampMs: Date.parse("2026-08-07T02:40:13.944Z") }),
200 record({ timestampMs: Date.parse("2026-08-07T03:40:13.944Z") }),
201 ],
202 "America/Los_Angeles",
203 "hour",
204 );
205​
206 expect(result.buckets.map((bucket) => [bucket.day, bucket.hourStart])).toEqual([
207 ["2026-08-06", "2026-08-07T02:37:00.000Z"],
208 ["2026-08-06", "2026-08-07T03:37:00.000Z"],
209 ]);
210 });
211​
212 it("uses an inclusive start and exclusive end for rolling windows", () => {
213 const result = aggregate(
214 [
215 record({ timestampMs: Date.parse("2026-08-06T04:36:59.999Z") }),
216 record({ timestampMs: Date.parse("2026-08-06T04:37:00.000Z") }),
217 record({ timestampMs: Date.parse("2026-08-07T04:36:59.999Z") }),
218 record({ timestampMs: Date.parse("2026-08-07T04:37:00.000Z") }),
219 ],
220 "UTC",
221 "hour",
222 );
223​
224 expect(result.outOfWindow).toBe(2);
225 expect(result.buckets.map((bucket) => bucket.hourStart)).toEqual([
226 "2026-08-06T04:37:00.000Z",
227 "2026-08-07T03:37:00.000Z",
228 ]);
229 });
230​
231 it("counts tokens but not cost for a model with no rate", () => {
232 const result = aggregate([record({ model: "kimi-k3" })]);
233​
234 expect(result.buckets[0]?.costUsd).toBe(0);
235 expect(result.buckets[0]?.costSource).toBe("unpriced");
236 expect(result.buckets[0]?.unpricedRecords).toBe(1);
237 expect(result.buckets[0]?.totals.outputTokens).toBe(50);
238 });
239​
240 it("prefers a reported cost over the rate table", () => {
241 const result = aggregate([record({ reportedCostUsd: 1.25 })]);
242​
243 expect(result.buckets[0]?.costUsd).toBe(1.25);
244 expect(result.buckets[0]?.costSource).toBe("providerReported");
245 });
246​
247 it("folds a mapped model into its target and prices it there", () => {
248 const aggregator = new UsageAggregator({
249 timeZone: "UTC",
250 sinceDay: "2026-08-01",
251 untilDay: "2026-08-31",
252 rates,
253 modelAliases: resolveModelAliases({ "example-preview": "claude-fable-5" }),
254 });
255 aggregator.add(record());
256 aggregator.add(record({ model: "example-preview", rateModel: "example-preview-high" }));
257 const result = aggregator.finish();
258​
259 expect(result.buckets).toHaveLength(1);
260 expect(result.buckets[0]?.model).toBe("claude-fable-5");
261 expect(result.buckets[0]?.records).toBe(2);
262 expect(result.buckets[0]?.costUsd).toBeCloseTo(0.00925, 9);
263 expect(result.buckets[0]?.unpricedRecords).toBe(0);
264 });
265});
266​
267describe("resolveModelAliases", () => {
268 it("follows chains to the final model and drops chains that enter a loop", () => {
269 expect(
270 resolveModelAliases({
271 "preview[1m]": "preview",
272 preview: "example-model",
273 loop: "back",
274 back: "loop",
275 intoLoop: "loop",
276 self: "self",
277 }),
278 ).toEqual(
279 new Map([
280 ["preview[1m]", "example-model"],
281 ["preview", "example-model"],
282 ]),
283 );
284 });
285});