| 1 | /** |
| 2 | * Search embeddings, as billing sees them. Every passage embedded is |
| 3 | * counted in `doc_embed_usage` (src/indexer.ts, tokens estimated as |
| 4 | * characters / 4). Every quarter hour (wrangler.jsonc `triggers`) each |
| 5 | * workspace's month so far goes to billing (`note_pending`, source |
| 6 | * `context`, the price book's `embedding_tokens`) at Workers AI's price: |
| 7 | * billing charges the plan once the month is over and covers a free |
| 8 | * workspace's itself (services/billing storage.rs `COVERED_FOR_FREE`). |
| 9 | * |
| 10 | * So that g1t's cover has an end, a workspace without the plan embeds at |
| 11 | * most `FREE_EMBED_TOKENS_PER_MONTH` a month; past it, its passages are |
| 12 | * kept for words, and embedded once a month starts or it takes the plan. |
| 13 | * The plan is the one billing last said (`workspace_billing.plan`); a |
| 14 | * workspace billing has said nothing about yet is held like a free one. |
| 15 | * |
| 16 | * Billing knows workspaces by slug and this service keeps usage by id, so |
| 17 | * each workspace's slug is noted when the service looks it up. |
| 18 | */ |
| 19 | import type { ServiceBinding } from "@g1t/contracts"; |
| 20 | |
| 21 | /** Embedding tokens a workspace without the plan may use a month: about 20 MB of text. */ |
| 22 | export const FREE_EMBED_TOKENS_PER_MONTH = 5_000_000; |
| 23 | |
| 24 | /** What Workers AI charges g1t a million embedding tokens, when the price book cannot be read ($0.067). */ |
| 25 | export const EMBEDDING_MICROS_PER_MILLION_TOKENS = 67_000; |
| 26 | |
| 27 | /** Workspaces whose slug this isolate has noted already. */ |
| 28 | const noted = new Set<string>(); |
| 29 | |
| 30 | /** `2026-10`. */ |
| 31 | export const monthOf = (now: Date) => now.toISOString().slice(0, 7); |
| 32 | |
| 33 | /** Notes which slug a workspace id is, once per isolate. Never throws. */ |
| 34 | export async function rememberWorkspace(db: D1Database, workspace: { id: string; slug: string } | null): Promise<void> { |
| 35 | if (!workspace?.id || !workspace.slug) return; |
| 36 | const key = `${workspace.id}:${workspace.slug}`; |
| 37 | if (noted.has(key)) return; |
| 38 | try { |
| 39 | await db |
| 40 | .prepare( |
| 41 | `INSERT INTO workspace_billing (workspace_id, slug) VALUES (?1, ?2) |
| 42 | ON CONFLICT (workspace_id) DO UPDATE SET slug = excluded.slug WHERE slug <> excluded.slug`, |
| 43 | ) |
| 44 | .bind(workspace.id, workspace.slug.toLowerCase()) |
| 45 | .run(); |
| 46 | noted.add(key); |
| 47 | } catch (error) { |
| 48 | console.error("could not note the workspace's slug", workspace.id, String(error)); |
| 49 | } |
| 50 | } |
| 51 | |
| 52 | /** Embedding tokens the workspace used in `month`. */ |
| 53 | export async function monthTokens(db: D1Database, workspaceId: string, month: string): Promise<number> { |
| 54 | const row = await db |
| 55 | .prepare("SELECT COALESCE(SUM(tokens), 0) AS tokens FROM doc_embed_usage WHERE workspace_id = ? AND hour >= ? AND hour < ?") |
| 56 | .bind(workspaceId, month, `${month}~`) |
| 57 | .first<{ tokens: number }>(); |
| 58 | return Number(row?.tokens ?? 0); |
| 59 | } |
| 60 | |
| 61 | /** |
| 62 | * Whether the workspace has used its free month: only without the plan |
| 63 | * (or before billing has said which it is on). |
| 64 | */ |
| 65 | export async function overFreeBudget(db: D1Database, workspaceId: string, now: Date, budget = FREE_EMBED_TOKENS_PER_MONTH): Promise<boolean> { |
| 66 | const plan = await db.prepare("SELECT plan FROM workspace_billing WHERE workspace_id = ?").bind(workspaceId).first<{ plan: string | null }>(); |
| 67 | if (plan?.plan === "paid") return false; |
| 68 | return (await monthTokens(db, workspaceId, monthOf(now))) >= budget; |
| 69 | } |
| 70 | |
| 71 | /** What `tokens` cost g1t at `microsPerMillion`, rounded up. */ |
| 72 | export function embeddingCost(tokens: number, microsPerMillion: number): number { |
| 73 | return Math.ceil((Math.max(0, tokens) / 1_000_000) * microsPerMillion); |
| 74 | } |
| 75 | |
| 76 | /** For the Billing page: `1.2 million tokens`. */ |
| 77 | export function embeddingDetail(tokens: number): string { |
| 78 | return tokens >= 1_000_000 ? `${(tokens / 1_000_000).toFixed(1)} million tokens` : `${tokens.toLocaleString("en-US")} tokens`; |
| 79 | } |
| 80 | |
| 81 | type Billing = { |
| 82 | freeWorkspaces(workspaces: string[]): Promise<string[]>; |
| 83 | prices(): Promise<{ prices: { meter: string; costMicros: number }[] }>; |
| 84 | notePending(workspace: string, source: "context", costMicros: number, detail?: string | null, month?: string | null): Promise<boolean>; |
| 85 | }; |
| 86 | |
| 87 | /** The month before `month` (`2026-10` → `2026-09`). */ |
| 88 | export function previousMonth(month: string): string { |
| 89 | const [year, number] = month.split("-").map(Number); |
| 90 | return new Date(Date.UTC(year, number - 2, 1)).toISOString().slice(0, 7); |
| 91 | } |
| 92 | |
| 93 | type MonthRow = { workspace_id: string; tokens: number; slug: string | null; noted_month: string | null; noted_tokens: number | null; month: string }; |
| 94 | |
| 95 | /** |
| 96 | * Reports each workspace's month of embedding tokens that changed since |
| 97 | * it was last reported, naming the month the total is for, and learns |
| 98 | * which of them are on the plan. A workspace whose slug is not known yet |
| 99 | * waits for the next run. Once a month is over, its last tokens still go |
| 100 | * as that month's for a workspace last reported in it (billing takes the |
| 101 | * month before until it charges it). Never throws. |
| 102 | */ |
| 103 | export async function reportEmbeddings(db: D1Database, billing: Billing, now = new Date()): Promise<{ reported: number; unknown: number }> { |
| 104 | const month = monthOf(now); |
| 105 | const previous = previousMonth(month); |
| 106 | const totals = (of: string) => |
| 107 | db |
| 108 | .prepare( |
| 109 | `SELECT u.workspace_id AS workspace_id, SUM(u.tokens) AS tokens, b.slug AS slug, b.noted_month AS noted_month, b.noted_tokens AS noted_tokens, ?1 AS month |
| 110 | FROM doc_embed_usage u LEFT JOIN workspace_billing b ON b.workspace_id = u.workspace_id |
| 111 | WHERE u.hour >= ?1 AND u.hour < ?2 |
| 112 | GROUP BY u.workspace_id`, |
| 113 | ) |
| 114 | .bind(of, `${of}~`) |
| 115 | .all<MonthRow>(); |
| 116 | const [last, rows] = await Promise.all([totals(previous), totals(month)]); |
| 117 | // Last month's end: only for workspaces last reported in it, so a month |
| 118 | // reported in full (or never) is not sent again. |
| 119 | const ended = last.results.filter((r) => r.slug && r.noted_month === previous && Number(r.noted_tokens ?? 0) !== Number(r.tokens ?? 0)); |
| 120 | const known = rows.results.filter((r) => r.slug); |
| 121 | const unknown = rows.results.length - known.length; |
| 122 | if (unknown) console.log("embedding usage of workspaces not looked up yet waits", unknown); |
| 123 | if (!known.length) return { reported: 0, unknown }; |
| 124 | // Which are free, so the budget holds; when billing does not answer, the plan it last said. |
| 125 | try { |
| 126 | const free = new Set((await billing.freeWorkspaces(known.map((r) => r.slug!))).map((s) => s.toLowerCase())); |
| 127 | const at = now.toISOString(); |
| 128 | const updates = known.map((r) => |
| 129 | db.prepare("UPDATE workspace_billing SET plan = ?, plan_at = ? WHERE workspace_id = ?").bind(free.has(r.slug!) ? "free" : "paid", at, r.workspace_id), |
| 130 | ); |
| 131 | for (let i = 0; i < updates.length; i += 50) await db.batch(updates.slice(i, i + 50)); |
| 132 | } catch (error) { |
| 133 | console.error("could not tell which workspaces are free; keeping what billing last said", String(error)); |
| 134 | } |
| 135 | const book = await billing.prices().catch(() => null); |
| 136 | const perMillion = book?.prices.find((p) => p.meter === "embedding_tokens")?.costMicros ?? EMBEDDING_MICROS_PER_MILLION_TOKENS; |
| 137 | let reported = 0; |
| 138 | for (const r of [...ended, ...known]) { |
| 139 | const tokens = Number(r.tokens ?? 0); |
| 140 | if (r.noted_month === r.month && Number(r.noted_tokens ?? 0) === tokens) continue; |
| 141 | try { |
| 142 | // Any answer is billing's to keep (a month it has charged is refused |
| 143 | // for good); only a failure to reach it is tried again. |
| 144 | await billing.notePending(r.slug!, "context", embeddingCost(tokens, perMillion), embeddingDetail(tokens), r.month); |
| 145 | await db |
| 146 | .prepare("UPDATE workspace_billing SET noted_month = ?, noted_tokens = ?, noted_at = ? WHERE workspace_id = ?") |
| 147 | .bind(r.month, tokens, now.toISOString(), r.workspace_id) |
| 148 | .run(); |
| 149 | reported++; |
| 150 | } catch (error) { |
| 151 | // Tried again on the next run. |
| 152 | console.error("embedding usage not reported", r.slug, tokens, String(error)); |
| 153 | } |
| 154 | } |
| 155 | return { reported, unknown }; |
| 156 | } |
| 157 | |
| 158 | /** Billing, as `reportEmbeddings` uses it. */ |
| 159 | export async function billingFor(service: ServiceBinding): Promise<Billing> { |
| 160 | // Loaded here, so the rules above test under Node without the contracts' runtime. |
| 161 | const { billingClient } = await import("@g1t/contracts"); |
| 162 | return billingClient(service) as unknown as Billing; |
| 163 | } |