Skip to content

Where people and agents ship software together. The open-source git platform for the whole job: issues, agents, checks and deploys to the edge.

g1t.sh
146 linesCodeBlameRaw
1import assert from "node:assert/strict";
2import { readFileSync, readdirSync } from "node:fs";
3import { DatabaseSync } from "node:sqlite";
4import { test } from "node:test";
5
6import {
7 FREE_EMBED_TOKENS_PER_MONTH,
8 embeddingCost,
9 embeddingDetail,
10 monthTokens,
11 overFreeBudget,
12 rememberWorkspace,
13 reportEmbeddings,
14} from "./embed-billing.ts";
15import { indexDoc } from "./indexer.ts";
16import type { Embedder, VectorStore } from "./vectors.ts";
17
18function 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
37const NOW = new Date("2026-10-13T12:00:00Z");
38
39function 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
43function 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
64test("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
84test("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
104test("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
119test("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
141test("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});