1// @effect-diagnostics-next-line nodeBuiltinImport:off - Effect's File.Info has no ctime, which the cache fingerprint needs.
2import * as NodeFSP from "node:fs/promises";
3
4import type { UsageRecord } from "@t3tools/provider-core/server/usage";
5import * as NodeSqliteClient from "@t3tools/shared/nodeSqliteClient";
6import * as Effect from "effect/Effect";
7import * as Exit from "effect/Exit";
8import * as FileSystem from "effect/FileSystem";
9import * as Path from "effect/Path";
10import * as Schema from "effect/Schema";
11import * as SqlClient from "effect/sql/SqlClient";
12
13type FieldValue = number | bigint | Uint8Array;
14type Fields = Map<number, FieldValue[]>;
15
16/** Antigravity stores usage metadata as protobuf, independently of conversation text. */
17function fields(bytes: Uint8Array): Fields {
18 let offset = 0;
19 const result: Fields = new Map();
20 const varint = () => {
21 let value = 0n;
22 for (let shift = 0n; shift < 70n; shift += 7n) {
23 const byte = bytes[offset++];
24 if (byte === undefined || (shift === 63n && byte > 1)) {
25 throw new Error("Invalid Antigravity protobuf varint");
26 }
27 value |= BigInt(byte & 127) << shift;
28 if (byte < 128) {
29 return value > BigInt(Number.MAX_SAFE_INTEGER) ? value : Number(value);
30 }
31 }
32 throw new Error("Invalid Antigravity protobuf varint");
33 };
34 while (offset < bytes.length) {
35 const tag = varint();
36 if (typeof tag !== "number") throw new Error("Invalid protobuf field");
37 const number = Math.floor(tag / 8);
38 const wire = tag % 8;
39 if (number === 0) throw new Error("Invalid protobuf field");
40 let value: FieldValue;
41 if (wire === 0) {
42 value = varint();
43 } else if (wire === 1 || wire === 5 || wire === 2) {
44 const length = wire === 2 ? varint() : wire === 1 ? 8 : 4;
45 if (typeof length !== "number") throw new Error("Invalid protobuf field length");
46 if (length > bytes.length - offset) throw new Error("Truncated protobuf field");
47 value = bytes.subarray(offset, offset + length);
48 offset += length;
49 if (wire !== 2) continue;
50 } else {
51 throw new Error("Unsupported protobuf wire type");
52 }
53 const entries = result.get(number) ?? [];
54 entries.push(value);
55 result.set(number, entries);
56 }
57 return result;
58}
59
60const numberAt = (value: Fields, key: number) => {
61 const entry = value.get(key)?.[0];
62 return typeof entry === "number" ? entry : 0;
63};
64const bytesAt = (value: Fields, key: number) => {
65 const entry = value.get(key)?.[0];
66 return entry instanceof Uint8Array ? entry : undefined;
67};
68const nested = (value: Fields, key: number) => {
69 const bytes = bytesAt(value, key);
70 return bytes === undefined ? new Map<number, FieldValue[]>() : fields(bytes);
71};
72const textAt = (value: Fields, key: number) => {
73 const bytes = bytesAt(value, key);
74 return bytes === undefined ? "" : new TextDecoder("utf-8", { fatal: true }).decode(bytes).trim();
75};
76const timestamp = (value: Fields) => {
77 const seconds = numberAt(value, 1);
78 return seconds > 0 ? seconds * 1000 + Math.floor(numberAt(value, 2) / 1_000_000) : null;
79};
80
81const MODEL_IDS: Record<number, string> = {
82 246: "gemini-2.5-pro",
83 312: "gemini-2.5-flash",
84 313: "gemini-2.5-flash-thinking",
85 329: "gemini-2.5-flash-thinking",
86 330: "gemini-2.5-flash-lite",
87 281: "claude-sonnet-4",
88 282: "claude-sonnet-4",
89 290: "claude-opus-4",
90 291: "claude-opus-4",
91 333: "claude-sonnet-4-5",
92 334: "claude-sonnet-4-5",
93 340: "claude-haiku-4-5",
94 341: "claude-haiku-4-5",
95 1026: "claude-opus-4-6",
96 1035: "claude-sonnet-4-6",
97 1016: "gemini-3.1-pro",
98 1036: "gemini-3.1-pro",
99 1037: "gemini-3.1-pro",
100 1018: "gemini-3-flash-preview",
101 1084: "gemini-3-flash-preview",
102 1047: "gemini-3-flash-preview",
103};
104
105function modelName(name: string, id: number): string {
106 if (name) {
107 const normalized = name
108 .toLowerCase()
109 .replace(/\s*\([^)]*\)\s*$/, "")
110 .replaceAll(" ", "-");
111 if (normalized.startsWith("claude-")) {
112 return normalized
113 .replace(/^claude-(4(?:\.\d+)?)-(sonnet|opus|haiku)/, "claude-$2-$1")
114 .replaceAll(".", "-");
115 }
116 return normalized;
117 }
118 return MODEL_IDS[id] ?? (id > 0 ? `antigravity-model-${id}` : "");
119}
120
121interface Metadata {
122 model: string;
123 timestampMs: number | null;
124 usages: Fields[];
125}
126
127function metadata(bytes: Uint8Array, step: boolean): Metadata {
128 const root = fields(bytes);
129 if (!step && bytesAt(root, 1) === undefined) {
130 throw new Error("Missing Antigravity generation metadata");
131 }
132 const data = step ? root : nested(root, 1);
133 const model = step ? nested(data, 24) : data;
134 const usage = bytesAt(data, step ? 9 : 4);
135 // Copy the usage bytes: subarrays kept in usages would hold each row's whole blob alive.
136 const usages = usage === undefined ? [] : [fields(usage.slice())];
137 for (const retry of data.get(step ? 28 : 17) ?? []) {
138 if (!(retry instanceof Uint8Array)) throw new Error("Invalid retry metadata");
139 const retryUsage = bytesAt(fields(retry), 2);
140 if (retryUsage !== undefined) usages.push(fields(retryUsage.slice()));
141 }
142 return {
143 model: modelName(
144 textAt(model, step ? 12 : 19) || textAt(model, step ? 8 : 21),
145 numberAt(model, step ? 1 : 3),
146 ),
147 timestampMs: step
148 ? (timestamp(nested(data, 8)) ?? timestamp(nested(data, 1)))
149 : timestamp(nested(nested(data, 9), 4)),
150 usages,
151 };
152}
153
154function blob(value: unknown): Uint8Array {
155 if (!(value instanceof Uint8Array)) throw new Error("Invalid Antigravity metadata blob");
156 return value;
157}
158
159interface UsageCandidate {
160 record: UsageRecord;
161 keys: readonly string[];
162 timestampQuality: number;
163}
164
165/** A conversation database whose usage tables are missing or undecodable. */
166class AntigravityDatabaseDecodeError extends Schema.TaggedError<AntigravityDatabaseDecodeError>()(
167 "AntigravityDatabaseDecodeError",
168 { path: Schema.String, cause: Schema.optional(Schema.Defect()) },
169) {
170 override get message(): string {
171 return `Antigravity usage could not be decoded from ${this.path}.`;
172 }
173}
174
175/**
176 * Rows per statement. One statement.all() holds the event loop until every row
177 * it returns is in memory, and a generation row carries the conversation's
178 * whole context so far, often megabytes, so generations are read one at a time.
179 */
180const GENERATION_PAGE_ROWS = 1;
181const STEP_PAGE_ROWS = 64;
182
183const readDatabase = Effect.fn("readDatabase")(function* (
184 path: string,
185 sessionId: string,
186 fallbackTimestamp: number,
187) {
188 const sql = yield* SqlClient.SqlClient;
189 const decode = <A>(evaluate: () => A) =>
190 Effect.try({
191 try: evaluate,
192 catch: (cause) => new AntigravityDatabaseDecodeError({ path, cause }),
193 });
194 yield* sql.unsafe("PRAGMA busy_timeout = 100");
195 return yield* Effect.gen(function* () {
196 const tables = new Set(
197 (yield* sql.unsafe<{ readonly name: unknown }>(
198 "SELECT name FROM sqlite_master WHERE type = 'table'",
199 )).map((row) => row.name),
200 );
201 if (!tables.has("gen_metadata") && !tables.has("steps")) {
202 return yield* new AntigravityDatabaseDecodeError({ path });
203 }
204 // Runs `query` a page at a time, binding the last rowid read and the page
205 // size. A rowid seek works whatever the table's indexes. Antigravity declares
206 // idx as the INTEGER PRIMARY KEY, so rowid order is idx order; the sort keeps
207 // idx order where rows are stored otherwise.
208 const readMetadata = Effect.fnUntraced(function* (
209 query: string,
210 pageRows: number,
211 step: boolean,
212 ) {
213 const entries: Array<{ idx: number; entry: Metadata }> = [];
214 // Below every rowid.
215 let after = -Infinity;
216 while (true) {
217 // Untraced: a cold read runs one of these per generation.
218 const rows = yield* sql
219 .unsafe<{
220 readonly row_id: number;
221 readonly idx: unknown;
222 readonly data: unknown;
223 }>(query, [after, pageRows])
224 .pipe(Effect.withTracerEnabled(false));
225 for (const row of rows) {
226 // Skipped here rather than in SQL, so a statement never scans past
227 // more than one page of steps that have no metadata.
228 if (step && row.data === null) continue;
229 const idx = row.idx;
230 if (typeof idx !== "number") {
231 return yield* new AntigravityDatabaseDecodeError({ path });
232 }
233 entries.push({ idx, entry: yield* decode(() => metadata(blob(row.data), step)) });
234 }
235 if (rows.length < pageRows) break;
236 after = rows[rows.length - 1]!.row_id;
237 yield* Effect.yieldNow;
238 }
239 return entries.sort((a, b) => a.idx - b.idx);
240 });
241 const generations = tables.has("gen_metadata")
242 ? yield* readMetadata(
243 "SELECT rowid AS row_id, idx, data FROM gen_metadata WHERE rowid > ? ORDER BY rowid LIMIT ?",
244 GENERATION_PAGE_ROWS,
245 false,
246 )
247 : [];
248 let trajectoryTimestamp: number | null = null;
249 if (tables.has("trajectory_metadata_blob")) {
250 for (const row of yield* sql.unsafe<{ readonly data: unknown }>(
251 "SELECT data FROM trajectory_metadata_blob",
252 )) {
253 trajectoryTimestamp ??= yield* decode(() => timestamp(nested(fields(blob(row.data)), 2)));
254 }
255 }
256 const steps = tables.has("steps")
257 ? yield* readMetadata(
258 "SELECT rowid AS row_id, idx, metadata AS data FROM steps WHERE rowid > ? ORDER BY rowid LIMIT ?",
259 STEP_PAGE_ROWS,
260 true,
261 )
262 : [];
263 return yield* decode(() =>
264 usageCandidates(sessionId, generations, steps, trajectoryTimestamp, fallbackTimestamp),
265 );
266 }).pipe(sql.withTransaction);
267});
268
269function usageCandidates(
270 sessionId: string,
271 generations: ReadonlyArray<{ idx: number; entry: Metadata }>,
272 steps: ReadonlyArray<{ idx: number; entry: Metadata }>,
273 trajectoryTimestamp: number | null,
274 fallbackTimestamp: number,
275): UsageCandidate[] {
276 const records: UsageCandidate[] = [];
277 const generationModels = new Map(generations.map(({ idx, entry }) => [idx, entry.model]));
278 for (const [source, entries] of [
279 ["step", steps],
280 ["generation", generations],
281 ] as const) {
282 for (const [index, { idx, entry }] of entries.entries()) {
283 for (const [usageIndex, usage] of entry.usages.entries()) {
284 const outputTokens = Math.max(numberAt(usage, 3), numberAt(usage, 9) + numberAt(usage, 10));
285 const totals = {
286 uncachedInputTokens: numberAt(usage, 2),
287 cachedInputTokens: numberAt(usage, 5),
288 cacheCreationTokens: numberAt(usage, 4),
289 outputTokens,
290 reasoningTokens: Math.min(outputTokens, numberAt(usage, 9)),
291 };
292 if (
293 totals.uncachedInputTokens +
294 totals.cachedInputTokens +
295 totals.cacheCreationTokens +
296 outputTokens ===
297 0
298 )
299 continue;
300 const keys = ([11, 12, 7] as const).flatMap((key) => {
301 const id = textAt(usage, key);
302 return id ? [`antigravity:${key}:${id}`] : [];
303 });
304 const record: UsageRecord = {
305 provider: "antigravity",
306 sessionId,
307 timestampMs: entry.timestampMs ?? trajectoryTimestamp ?? fallbackTimestamp,
308 model:
309 MODEL_IDS[numberAt(usage, 1)] ||
310 entry.model ||
311 (source === "step" ? generationModels.get(idx) : "") ||
312 modelName("", numberAt(usage, 1)) ||
313 "antigravity-unknown",
314 totals,
315 reportedCostUsd: null,
316 speed: "standard",
317 dedupeKey: keys[0] ?? `antigravity:${sessionId}:${source}:${index}:${usageIndex}`,
318 };
319 records.push({
320 record,
321 keys,
322 timestampQuality: entry.timestampMs !== null ? 2 : trajectoryTimestamp !== null ? 1 : 0,
323 });
324 }
325 }
326 }
327 return records;
328}
329
330interface CachedDatabase {
331 readonly fingerprint: string;
332 readonly candidates: readonly UsageCandidate[];
333}
334
335/**
336 * Parsed databases keyed by canonical path. An entry is reused only while both
337 * the database and its `-wal` sidecar keep the same size, mtime and ctime: new
338 * rows land in the WAL without touching the main file until a checkpoint, and
339 * ctime moves on any content write even when mtime is restored.
340 */
341export const makeAntigravityUsageCache = () => new Map<string, CachedDatabase>();
342
343/**
344 * Reads and merges aliases across every configured Antigravity store before date
345 * filtering. With a cache, databases unchanged since the previous read are not
346 * decoded again.
347 */
348export const readAntigravityUsage = Effect.fn("readAntigravityUsage")(function* (
349 conversationsDirectories: string | readonly string[],
350 sinceMs: number,
351 cache?: Map<string, CachedDatabase>,
352) {
353 const fileSystem = yield* FileSystem.FileSystem;
354 const pathService = yield* Path.Path;
355 const roots =
356 typeof conversationsDirectories === "string"
357 ? [conversationsDirectories]
358 : conversationsDirectories;
359 const files: Array<{ root: string; path: string; records: UsageRecord[] }> = [];
360 const errors: string[] = [];
361 const identities = new Map<string, number>();
362 const groups: Array<
363 UsageCandidate & { parent: number; size: number; owner: number; fileIndex: number }
364 > = [];
365 const find = (index: number): number => {
366 let root = index;
367 while (groups[root]!.parent !== root) root = groups[root]!.parent;
368 while (index !== root) {
369 const parent = groups[index]!.parent;
370 groups[index]!.parent = root;
371 index = parent;
372 }
373 return root;
374 };
375 const merge = (left: number, right: number): number => {
376 let a = find(left);
377 let b = find(right);
378 if (a === b) return a;
379 if (groups[a]!.size < groups[b]!.size) [a, b] = [b, a];
380 const target = groups[a]!;
381 const source = groups[b]!;
382 const first = target.owner < source.owner ? target : source;
383 const bestTime =
384 source.timestampQuality > target.timestampQuality ||
385 (source.timestampQuality === target.timestampQuality &&
386 source.record.timestampMs < target.record.timestampMs)
387 ? source
388 : target;
389 const x = target.record.totals;
390 const y = source.record.totals;
391 target.record = {
392 ...first.record,
393 model:
394 first.record.model === "antigravity-unknown"
395 ? first === target
396 ? source.record.model
397 : target.record.model
398 : first.record.model,
399 timestampMs: bestTime.record.timestampMs,
400 totals: {
401 uncachedInputTokens: Math.max(x.uncachedInputTokens, y.uncachedInputTokens),
402 cachedInputTokens: Math.max(x.cachedInputTokens, y.cachedInputTokens),
403 cacheCreationTokens: Math.max(x.cacheCreationTokens, y.cacheCreationTokens),
404 outputTokens: Math.max(x.outputTokens, y.outputTokens),
405 reasoningTokens: Math.max(x.reasoningTokens, y.reasoningTokens),
406 },
407 };
408 target.timestampQuality = bestTime.timestampQuality;
409 target.owner = first.owner;
410 target.fileIndex = first.fileIndex;
411 target.size += source.size;
412 source.parent = a;
413 return a;
414 };
415 const append = (candidate: UsageCandidate, fileIndex: number) => {
416 const index = groups.length;
417 groups.push({ ...candidate, parent: index, size: 1, owner: index, fileIndex });
418 for (const key of candidate.keys) {
419 const existing = identities.get(key);
420 if (existing !== undefined) merge(index, existing);
421 identities.set(key, index);
422 }
423 };
424 const visited = new Set<string>();
425 const readFile = Effect.fn("readAntigravityUsage.readFile")(function* (path: string) {
426 const canonical = yield* fileSystem.realPath(path);
427 if (visited.has(canonical)) return null;
428 visited.add(canonical);
429 const fingerprintOf = (file: string) =>
430 Effect.tryPromise(() => NodeFSP.stat(file)).pipe(
431 Effect.map((stat) => ({ stat, part: `${stat.size}:${stat.mtimeMs}:${stat.ctimeMs}` })),
432 );
433 const { stat, part } = yield* fingerprintOf(path);
434 const wal = yield* fingerprintOf(`${path}-wal`).pipe(Effect.option);
435 const fingerprint = `${part}/${wal._tag === "Some" ? wal.value.part : "-"}`;
436 const cached = cache?.get(canonical);
437 if (cached?.fingerprint === fingerprint) return cached.candidates;
438 const candidates = yield* readDatabase(
439 path,
440 pathService.basename(path, ".db"),
441 stat.mtimeMs,
442 ).pipe(Effect.provide(NodeSqliteClient.layer({ filename: path, readonly: true })));
443 cache?.set(canonical, { fingerprint, candidates });
444 return candidates;
445 });
446 const walk = (directory: string, root: string): Effect.Effect<void> =>
447 Effect.gen(function* () {
448 const names = yield* fileSystem.readDirectory(directory).pipe(
449 Effect.catchTags({
450 PlatformError: (cause) => {
451 if (cause.reason._tag !== "NotFound") errors.push(directory);
452 return Effect.succeed(null);
453 },
454 }),
455 );
456 if (names === null) return;
457 names.sort((a, b) => a.localeCompare(b));
458 for (const name of names) {
459 const path = pathService.join(directory, name);
460 // Like a directory listing's entry types: symlinks are neither files nor directories.
461 const isLink = Exit.isSuccess(yield* Effect.exit(fileSystem.readLink(path)));
462 const type = isLink
463 ? "SymbolicLink"
464 : yield* fileSystem.stat(path).pipe(
465 Effect.map((info) => info.type),
466 Effect.orElseSucceed(() => null),
467 );
468 if (type === "Directory") {
469 yield* walk(path, root);
470 } else if ((type === "File" || type === null) && name.endsWith(".db")) {
471 const candidates = yield* readFile(path).pipe(
472 Effect.catch(() => {
473 errors.push(path);
474 return Effect.succeed(null);
475 }),
476 );
477 if (candidates === null) continue;
478 const fileIndex = files.length;
479 files.push({ root, path, records: [] });
480 for (const [index, candidate] of candidates.entries()) {
481 append(candidate, fileIndex);
482 if (index % 256 === 255) yield* Effect.yieldNow;
483 }
484 }
485 }
486 });
487 for (const root of roots) yield* walk(root, root);
488 if (cache !== undefined) {
489 for (const key of cache.keys()) if (!visited.has(key)) cache.delete(key);
490 }
491 for (const [index, group] of groups.entries()) {
492 if (group.parent === index && group.record.timestampMs >= sinceMs) {
493 files[group.fileIndex]!.records.push(group.record);
494 }
495 }
496 return { files, errors };
497});