| 1 | import assert from "node:assert/strict"; |
| 2 | import { readFileSync, readdirSync } from "node:fs"; |
| 3 | import { DatabaseSync } from "node:sqlite"; |
| 4 | import { test } from "node:test"; |
| 5 | |
| 6 | import { |
| 7 | FREE_EMBED_TOKENS_PER_MONTH, |
| 8 | embeddingCost, |
| 9 | embeddingDetail, |
| 10 | monthTokens, |
| 11 | overFreeBudget, |
| 12 | rememberWorkspace, |
| 13 | reportEmbeddings, |
| 14 | } from "./embed-billing.ts"; |
| 15 | import { indexDoc } from "./indexer.ts"; |
| 16 | import type { Embedder, VectorStore } from "./vectors.ts"; |
| 17 | |
| 18 | function fakeD1() { |
| 19 | const db = new DatabaseSync(":memory:"); |
| 20 | const dir = new URL("../migrations/", import.meta.url); |
| 21 | for (const file of readdirSync(dir).sort()) db.exec(readFileSync(new URL(file, dir), "utf8")); |
| 22 | const statement = (sql: string, params: unknown[] = []) => ({ |
| 23 | sql, |
| 24 | params, |
| 25 | bind: (...values: unknown[]) => statement(sql, values), |
| 26 | first: async () => (db.prepare(sql).get(...(params as never[])) as unknown) ?? null, |
| 27 | all: async () => ({ results: db.prepare(sql).all(...(params as never[])) }), |
| 28 | run: async () => db.prepare(sql).run(...(params as never[])), |
| 29 | }); |
| 30 | const d1 = { |
| 31 | prepare: (sql: string) => statement(sql), |
| 32 | batch: async (list: ReturnType<typeof statement>[]) => list.map((s) => db.prepare(s.sql).run(...(s.params as never[]))), |
| 33 | } as unknown as D1Database; |
| 34 | return { d1, db }; |
| 35 | } |
| 36 | |
| 37 | const NOW = new Date("2026-10-13T12:00:00Z"); |
| 38 | |
| 39 | function usage(db: DatabaseSync, workspace: string, hour: string, tokens: number) { |
| 40 | db.prepare("INSERT INTO doc_embed_usage (workspace_id, hour, chunks, tokens) VALUES (?, ?, 1, ?)").run(workspace, hour, tokens); |
| 41 | } |
| 42 | |
| 43 | function billing(free: string[] = [], fail = false) { |
| 44 | const notes: [string, string, number, string | null | undefined, string | null | undefined][] = []; |
| 45 | return { |
| 46 | notes, |
| 47 | client: { |
| 48 | async freeWorkspaces(slugs: string[]) { |
| 49 | if (fail) throw new Error("billing answered 503"); |
| 50 | return slugs.filter((s) => free.includes(s)); |
| 51 | }, |
| 52 | async prices() { |
| 53 | return { prices: [{ meter: "embedding_tokens", costMicros: 67_000 }] }; |
| 54 | }, |
| 55 | async notePending(workspace: string, source: "context", cost: number, detail?: string | null, month?: string | null) { |
| 56 | if (fail) throw new Error("billing answered 503"); |
| 57 | notes.push([workspace, source, cost, detail, month]); |
| 58 | return true; |
| 59 | }, |
| 60 | }, |
| 61 | }; |
| 62 | } |
| 63 | |
| 64 | test("each workspace's month of tokens goes to billing as context, once per change, and only this month's", async () => { |
| 65 | const { d1, db } = fakeD1(); |
| 66 | await rememberWorkspace(d1, { id: "wsp_1", slug: "Acme" }); |
| 67 | usage(db, "wsp_1", "2026-10-01T03", 1_500_000); |
| 68 | usage(db, "wsp_1", "2026-10-13T11", 500_000); |
| 69 | usage(db, "wsp_1", "2026-09-30T23", 9_000_000); |
| 70 | // Not looked up yet: waits. |
| 71 | usage(db, "wsp_2", "2026-10-13T11", 10); |
| 72 | assert.equal(await monthTokens(d1, "wsp_1", "2026-10"), 2_000_000); |
| 73 | const b = billing(["acme"]); |
| 74 | assert.deepEqual(await reportEmbeddings(d1, b.client, NOW), { reported: 1, unknown: 1 }); |
| 75 | assert.deepEqual(b.notes, [["acme", "context", 134_000, "2.0 million tokens", "2026-10"]]); |
| 76 | // Nothing new: nothing sent. |
| 77 | await reportEmbeddings(d1, b.client, NOW); |
| 78 | assert.equal(b.notes.length, 1); |
| 79 | usage(db, "wsp_1", "2026-10-13T12", 1000); |
| 80 | await reportEmbeddings(d1, b.client, NOW); |
| 81 | assert.deepEqual(b.notes.at(-1), ["acme", "context", 134_067, "2.0 million tokens", "2026-10"]); |
| 82 | }); |
| 83 | |
| 84 | test("a month's last tokens go as that month's once it is over, then the new month's on their own", async () => { |
| 85 | const { d1, db } = fakeD1(); |
| 86 | await rememberWorkspace(d1, { id: "wsp_1", slug: "acme" }); |
| 87 | usage(db, "wsp_1", "2026-10-31T20", 1_000_000); |
| 88 | const b = billing([]); |
| 89 | await reportEmbeddings(d1, b.client, new Date("2026-10-31T21:00:00Z")); |
| 90 | assert.deepEqual(b.notes, [["acme", "context", 67_000, "1.0 million tokens", "2026-10"]]); |
| 91 | // The last quarter hour of October, reported in November: as October's. |
| 92 | usage(db, "wsp_1", "2026-10-31T23", 1_000_000); |
| 93 | usage(db, "wsp_1", "2026-11-01T00", 10); |
| 94 | await reportEmbeddings(d1, b.client, new Date("2026-11-01T00:15:00Z")); |
| 95 | assert.deepEqual(b.notes.slice(1), [ |
| 96 | ["acme", "context", 134_000, "2.0 million tokens", "2026-10"], |
| 97 | ["acme", "context", 1, "10 tokens", "2026-11"], |
| 98 | ]); |
| 99 | // Nothing new: October is not sent again. |
| 100 | await reportEmbeddings(d1, b.client, new Date("2026-11-01T00:30:00Z")); |
| 101 | assert.equal(b.notes.length, 3); |
| 102 | }); |
| 103 | |
| 104 | test("billing down: nothing is marked reported, and the plan it last said holds", async () => { |
| 105 | const { d1, db } = fakeD1(); |
| 106 | await rememberWorkspace(d1, { id: "wsp_9", slug: "paid" }); |
| 107 | usage(db, "wsp_9", "2026-10-13T11", FREE_EMBED_TOKENS_PER_MONTH); |
| 108 | // Unknown is held like a free workspace. |
| 109 | assert.equal(await overFreeBudget(d1, "wsp_9", NOW), true); |
| 110 | await reportEmbeddings(d1, billing([]).client, NOW); |
| 111 | assert.equal(await overFreeBudget(d1, "wsp_9", NOW), false, "on the plan, no budget"); |
| 112 | const down = billing([], true); |
| 113 | await reportEmbeddings(d1, down.client, NOW); |
| 114 | assert.equal(await overFreeBudget(d1, "wsp_9", NOW), false, "still the plan billing last said"); |
| 115 | const row = db.prepare("SELECT noted_tokens FROM workspace_billing WHERE workspace_id = 'wsp_9'").get() as { noted_tokens: number }; |
| 116 | assert.equal(row.noted_tokens, FREE_EMBED_TOKENS_PER_MONTH); |
| 117 | }); |
| 118 | |
| 119 | test("a free workspace past its month's budget embeds nothing more; its passages are kept for words", async () => { |
| 120 | const { d1, db } = fakeD1(); |
| 121 | await rememberWorkspace(d1, { id: "w1", slug: "free" }); |
| 122 | db.prepare("UPDATE workspace_billing SET plan = 'free'").run(); |
| 123 | usage(db, "w1", "2026-10-02T00", FREE_EMBED_TOKENS_PER_MONTH); |
| 124 | let embedded = 0; |
| 125 | const embedder: Embedder = { |
| 126 | async embed(texts) { |
| 127 | embedded += texts.length; |
| 128 | return texts.map(() => [1, 1]); |
| 129 | }, |
| 130 | }; |
| 131 | const store: VectorStore = { upsert: async () => {}, get: async () => [], delete: async () => {}, query: async () => [] }; |
| 132 | const result = await indexDoc(d1, embedder, store, { kind: "page", doc_id: "pag_1", workspace_id: "w1", space_id: "s1", title: "Ops", markdown: "## Deploy\n\nShip it carefully." }, NOW); |
| 133 | assert.equal(embedded, 0); |
| 134 | assert.equal(result.capped, true); |
| 135 | assert.ok(result.chunks > 0); |
| 136 | // A new month starts over. |
| 137 | await indexDoc(d1, embedder, store, { kind: "page", doc_id: "pag_1", workspace_id: "w1", space_id: "s1", title: "Ops", markdown: "## Deploy\n\nShip it carefully." }, new Date("2026-11-01T00:00:00Z")); |
| 138 | assert.ok(embedded > 0); |
| 139 | }); |
| 140 | |
| 141 | test("cost and detail", () => { |
| 142 | assert.equal(embeddingCost(1_000_000, 67_000), 67_000); |
| 143 | assert.equal(embeddingCost(1, 67_000), 1); |
| 144 | assert.equal(embeddingCost(-5, 67_000), 0); |
| 145 | assert.equal(embeddingDetail(1234), "1,234 tokens"); |
| 146 | }); |