| 1 | import assert from "node:assert/strict"; |
| 2 | import { test } from "node:test"; |
| 3 | |
| 4 | import type { ModelUpstream } from "@g1t/contracts"; |
| 5 | |
| 6 | import { isAnswer, tokenReport } from "./report.ts"; |
| 7 | import { NO_TOKENS, measure } from "./usage.ts"; |
| 8 | |
| 9 | const upstream: ModelUpstream = { |
| 10 | route: "g1t", |
| 11 | api: "anthropic", |
| 12 | model: null, |
| 13 | official: false, |
| 14 | provider: "g1t", |
| 15 | workspace: "acme", |
| 16 | repo: "acme/web", |
| 17 | number: 7, |
| 18 | task: "implement", |
| 19 | session: "ms_abc", |
| 20 | tier: "large", |
| 21 | requestedBy: "ada", |
| 22 | baseUrl: null, |
| 23 | apiKey: null, |
| 24 | authHeader: null, |
| 25 | }; |
| 26 | |
| 27 | test("only messages are answers; counting tokens is not", () => { |
| 28 | assert.equal(isAnswer("/v1/messages"), true); |
| 29 | assert.equal(isAnswer("/v1/messages?beta=true"), true); |
| 30 | assert.equal(isAnswer("/v1/messages/count_tokens?beta=true"), false); |
| 31 | assert.equal(isAnswer("/v1/models"), false); |
| 32 | }); |
| 33 | |
| 34 | test("an answer's tokens are reported under its run and person", () => { |
| 35 | const report = tokenReport(upstream, "claude-opus-4", { input: 3, output: 9, cacheRead: 100, cacheWrite: 0 }); |
| 36 | assert.deepEqual(report, { |
| 37 | workspace: "acme", |
| 38 | session: "ms_abc", |
| 39 | person: "ada", |
| 40 | model: "claude-opus-4", |
| 41 | tier: "large", |
| 42 | input: 3, |
| 43 | output: 9, |
| 44 | cacheRead: 100, |
| 45 | cacheWrite: 0, |
| 46 | }); |
| 47 | // A route that names a model counts under it. |
| 48 | assert.equal(tokenReport({ ...upstream, model: "gpt-x", tier: null }, "other", { ...NO_TOKENS, output: 1 })?.model, "gpt-x"); |
| 49 | }); |
| 50 | |
| 51 | test("nothing used, or no session to count it under, is not reported", () => { |
| 52 | assert.equal(tokenReport(upstream, null, { ...NO_TOKENS }), null); |
| 53 | assert.equal(tokenReport({ ...upstream, session: "" }, null, { ...NO_TOKENS, input: 1 }), null); |
| 54 | }); |
| 55 | |
| 56 | test("the model that answered is read from the answer", async () => { |
| 57 | const body = `data: ${JSON.stringify({ type: "message_start", message: { model: "claude-x", usage: { input_tokens: 1 } } })}\n\n`; |
| 58 | const streamed = measure(new Response(body, { headers: { "content-type": "text/event-stream" } })); |
| 59 | await streamed.response.text(); |
| 60 | assert.equal(await streamed.model, "claude-x"); |
| 61 | const whole = measure(Response.json({ model: "claude-y", usage: { output_tokens: 2 } })); |
| 62 | await whole.response.text(); |
| 63 | assert.equal(await whole.model, "claude-y"); |
| 64 | }); |