diff --git a/apps/server/src/provider/Drivers/AntigravityDriver.ts b/apps/server/src/provider/Drivers/AntigravityDriver.ts index f030341a83c1d3797892486f64faf028a30a7903..c0812af2a6dfd4612826fa3f557d8726417d9566 100644 --- a/apps/server/src/provider/Drivers/AntigravityDriver.ts +++ b/apps/server/src/provider/Drivers/AntigravityDriver.ts @@ -26,6 +26,7 @@ import { } from "../../textGeneration/AntigravityTextGeneration.ts"; import { makeAntigravityAuth, type AntigravityAuth } from "../AntigravityAuth.ts"; import * as AntigravityInstallation from "../AntigravityInstallation.ts"; +import * as AntigravityUsage from "./AntigravityUsage.ts"; import { antigravityAuthConfigIssue, antigravityAuthLabel, @@ -80,11 +81,16 @@ export type AntigravityDriverEnv = | ProviderEventLoggers.ProviderEventLoggers; /** Each instance owns its Google profile. Executable releases are shared by the environment. */ -export const AntigravityDriver: ProviderDriver = { +export const AntigravityDriver: ProviderDriver< + AntigravitySettings, + AntigravityDriverEnv, + AntigravityUsage.AntigravityUsage +> = { driverKind: DRIVER, metadata: { displayName: "Antigravity", supportsMultipleInstances: true }, configSchema: AntigravitySettings, defaultConfig: () => decodeSettings({}), + usage: AntigravityUsage.antigravityUsageReader, create: ({ instanceId, displayName, accentColor, environment, enabled, config }) => Effect.gen(function* () { const crypto = yield* Crypto.Crypto; diff --git a/apps/server/src/provider/Drivers/AntigravityUsage.ts b/apps/server/src/provider/Drivers/AntigravityUsage.ts new file mode 100644 index 0000000000000000000000000000000000000000..d1958788ce2dca7f72009956219908ffd6c1a16e --- /dev/null +++ b/apps/server/src/provider/Drivers/AntigravityUsage.ts @@ -0,0 +1,140 @@ +/** + * AntigravityUsage - Antigravity's usage history as a usage source. + * + * Antigravity keeps conversations in SQLite databases under its data + * directories and each instance's profile. Parsed databases stay in memory + * while a database and its WAL keep the same `(size, mtime, ctime)`, so a + * warm scan decodes only what changed. + * + * @module provider/Drivers/AntigravityUsage + */ +import * as NodeOS from "node:os"; + +import { type AntigravitySettings, UsageReadError } from "@t3tools/contracts"; +import { expandHomePath } from "@t3tools/provider-core/server/pathExpansion"; +import * as ProviderHost from "@t3tools/provider-core/server/ProviderHost"; +import type { + ProviderUsageInstance, + ProviderUsageReader, + ProviderUsageScan, +} from "@t3tools/provider-core/server/usage"; +import { HostProcessEnvironment } from "@t3tools/shared/hostProcess"; +import * as Context from "effect/Context"; +import * as Crypto from "effect/Crypto"; +import * as Effect from "effect/Effect"; +import * as FileSystem from "effect/FileSystem"; +import * as Layer from "effect/Layer"; +import * as Path from "effect/Path"; + +import { resolveAntigravityInstanceDirectories } from "../antigravityAuthSupport.ts"; +import { makeAntigravityUsageCache, readAntigravityUsage } from "./antigravityUsageReader.ts"; + +export class AntigravityUsage extends Context.Service< + AntigravityUsage, + { + /** Every Antigravity store: the data directories and each configured instance's profile. */ + readonly scan: (input: { + readonly instances: ReadonlyArray>; + readonly windowStartMs: number; + }) => Effect.Effect, UsageReadError>; + } +>()("t3/provider/Drivers/AntigravityUsage") {} + +const make = Effect.gen(function* () { + const crypto = yield* Crypto.Crypto; + const fileSystem = yield* FileSystem.FileSystem; + const path = yield* Path.Path; + const host = yield* ProviderHost.ProviderHost; + const hostEnvironment = yield* HostProcessEnvironment; + const cache = makeAntigravityUsageCache(); + + /** `ANTIGRAVITY_DATA_DIR` (comma-separated) or the defaults, canonicalized. */ + const dataRoots = Effect.gen(function* () { + const home = NodeOS.homedir(); + const configured = hostEnvironment["ANTIGRAVITY_DATA_DIR"] + ?.split(",") + .map((value) => value.trim()) + .filter(Boolean); + const defaults = [ + ...["antigravity", "antigravity-cli", "antigravity-ide", "antigravity-backup"].map((name) => + path.join(home, ".gemini", name), + ), + path.join(home, ".config", "antigravity"), + ]; + const canonical = new Set(); + for (const root of configured?.length ? configured : defaults) { + const resolved = path.resolve(expandHomePath(root)); + canonical.add( + yield* fileSystem.realPath(resolved).pipe(Effect.orElseSucceed(() => resolved)), + ); + } + return [...canonical]; + }); + + const scan: AntigravityUsage["Service"]["scan"] = Effect.fn("AntigravityUsage.scan")(function* ({ + instances, + windowStartMs, + }) { + const roots = yield* dataRoots; + // Only configured instances have a profile; the implicit default never ran. + for (const { instanceId } of instances.filter((instance) => instance.configured)) { + const directories = yield* resolveAntigravityInstanceDirectories( + host.paths.stateDir, + instanceId, + ).pipe( + Effect.provideService(Crypto.Crypto, crypto), + Effect.provideService(Path.Path, path), + Effect.mapError( + (cause) => + new UsageReadError({ + reason: "scanFailed", + detail: "Antigravity profile directory could not be resolved.", + cause, + }), + ), + ); + roots.push(path.join(directories.profile, "antigravity-acp")); + } + const conversationDirs = new Set(); + for (const root of roots) { + const resolvedRoot = yield* fileSystem.realPath(root).pipe(Effect.orElseSucceed(() => root)); + const nested = path.join(resolvedRoot, "conversations"); + const dir = (yield* fileSystem + .exists(nested) + .pipe(Effect.catchCause(() => Effect.succeed(false)))) + ? nested + : resolvedRoot; + conversationDirs.add(yield* fileSystem.realPath(dir).pipe(Effect.orElseSucceed(() => dir))); + } + const result = yield* Effect.promise(() => + readAntigravityUsage([...conversationDirs], windowStartMs, cache), + ); + const scanned: ProviderUsageScan[] = []; + for (const dir of conversationDirs) { + const exists = yield* fileSystem + .exists(dir) + .pipe(Effect.catchCause(() => Effect.succeed(false))); + const failed = result.errors.some( + (error) => error === dir || error.startsWith(`${dir}${path.sep}`), + ); + scanned.push({ + dir, + files: !exists && !failed ? null : result.files.filter((file) => file.root === dir), + status: failed ? "partial" : "ok", + ...(failed ? { message: "Some Antigravity history could not be read." } : {}), + }); + } + return scanned; + }); + + return AntigravityUsage.of({ scan }); +}); + +export const layer = Layer.effect(AntigravityUsage, make); + +export const antigravityUsageReader: ProviderUsageReader = { + kind: "scan", + provider: "antigravity", + scan: ({ instances, windowStartMs }) => + AntigravityUsage.pipe(Effect.flatMap((usage) => usage.scan({ instances, windowStartMs }))), +}; diff --git a/apps/server/src/provider/Drivers/antigravityUsageReader.test.ts b/apps/server/src/provider/Drivers/antigravityUsageReader.test.ts new file mode 100644 index 0000000000000000000000000000000000000000..edafd4ee89a6945371efb784a2c198464e747ca5 --- /dev/null +++ b/apps/server/src/provider/Drivers/antigravityUsageReader.test.ts @@ -0,0 +1,391 @@ +// @effect-diagnostics nodeBuiltinImport:off - the reader under test opens real +// SQLite databases, mirroring the reader's own deliberate node:sqlite usage. +import * as NodeFSP from "node:fs/promises"; +import * as NodeOS from "node:os"; +import * as NodePath from "node:path"; +import * as NodeSqlite from "node:sqlite"; + +import { afterEach, assert, beforeEach, describe, it } from "@effect/vitest"; + +import { makeAntigravityUsageCache, readAntigravityUsage } from "./antigravityUsageReader.ts"; + +let dir: string; + +function protoNumber(field: number, value: number): number[] { + const varint = (number: number) => { + const bytes: number[] = []; + do { + const byte = number % 128; + number = Math.floor(number / 128); + bytes.push(byte + (number > 0 ? 128 : 0)); + } while (number > 0); + return bytes; + }; + return [...varint(field * 8), ...varint(value)]; +} + +function protoBytes(field: number, bytes: readonly number[]): number[] { + const encoded = protoNumber(field, bytes.length); + encoded[0] = encoded[0]! + 2; + return [...encoded, ...bytes]; +} + +function protoText(field: number, value: string): number[] { + return protoBytes(field, [...Buffer.from(value)]); +} + +beforeEach(async () => { + dir = await NodeFSP.mkdtemp(NodePath.join(NodeOS.tmpdir(), "usage-reader-test-")); +}); + +afterEach(async () => { + await NodeFSP.rm(dir, { recursive: true, force: true }); +}); + +describe("readAntigravityUsage", () => { + it("deduplicates Antigravity generation and step usage while preserving retry model and token buckets", async () => { + const db = new NodeSqlite.DatabaseSync(NodePath.join(dir, "session-1.db")); + const stamp = protoNumber(1, 1780000000); + const usage = [ + ...protoNumber(2, 100), + ...protoNumber(3, 40), + ...protoNumber(4, 5), + ...protoNumber(5, 20), + ...protoNumber(9, 10), + ...protoText(11, "response-1"), + ]; + const retry = [ + ...protoNumber(1, 1026), + ...protoNumber(2, 12), + ...protoNumber(3, 3), + ...protoText(11, "retry-1"), + ]; + const generation = protoBytes(1, [ + ...protoBytes(4, usage), + ...protoText(19, "Gemini 3 Pro"), + ...protoBytes(9, protoBytes(4, stamp)), + ]); + const step = [ + ...protoBytes(9, usage), + ...protoBytes(8, stamp), + ...protoBytes(28, protoBytes(2, retry)), + ]; + try { + db.exec( + "CREATE TABLE gen_metadata (idx INTEGER, data BLOB); CREATE TABLE steps (idx INTEGER, metadata BLOB)", + ); + db.prepare("INSERT INTO gen_metadata VALUES (?, ?)").run(0, new Uint8Array(generation)); + db.prepare("INSERT INTO steps VALUES (?, ?)").run(0, new Uint8Array(step)); + } finally { + db.close(); + } + const result = await readAntigravityUsage(dir, 0); + assert.deepStrictEqual(result.errors, []); + const records = result.files.flatMap((file) => file.records); + assert.strictEqual(records.length, 2); + const main = records.find((record) => record.model === "gemini-3-pro"); + assert.isDefined(main); + assert.strictEqual(main?.timestampMs, 1780000000000); + assert.strictEqual(main?.sessionId, "session-1"); + assert.deepStrictEqual(main?.totals, { + uncachedInputTokens: 100, + cachedInputTokens: 20, + cacheCreationTokens: 5, + outputTokens: 40, + reasoningTokens: 10, + }); + assert.strictEqual( + records.find((record) => record.model === "claude-opus-4-6")?.totals.uncachedInputTokens, + 12, + ); + assert.deepStrictEqual( + (await readAntigravityUsage(dir, 1780000000001)).files.flatMap((file) => file.records), + [], + ); + }); + + it("uses the matching Antigravity generation model for each model-less step", async () => { + const db = new NodeSqlite.DatabaseSync(NodePath.join(dir, "model-switch.db")); + try { + db.exec( + "CREATE TABLE gen_metadata (idx INTEGER, data BLOB); CREATE TABLE steps (idx INTEGER, metadata BLOB)", + ); + const generation = db.prepare("INSERT INTO gen_metadata VALUES (?, ?)"); + const step = db.prepare("INSERT INTO steps VALUES (?, ?)"); + for (const [idx, name] of ["Gemini 3 Pro", "Claude Opus 4.6"].entries()) { + generation.run(idx, new Uint8Array(protoBytes(1, protoText(19, name)))); + step.run(idx, new Uint8Array(protoBytes(9, protoNumber(2, 10 + idx)))); + } + } finally { + db.close(); + } + const result = await readAntigravityUsage(dir, 0); + assert.deepStrictEqual(result.errors, []); + assert.deepStrictEqual( + result.files.flatMap((file) => file.records).map((record) => record.model), + ["gemini-3-pro", "claude-opus-4-6"], + ); + }); + + it("merges Antigravity aliases that bridge previously separate step records", async () => { + const db = new NodeSqlite.DatabaseSync(NodePath.join(dir, "bridge.db")); + try { + db.exec( + "CREATE TABLE gen_metadata (idx INTEGER, data BLOB); CREATE TABLE steps (idx INTEGER, metadata BLOB)", + ); + const step = db.prepare("INSERT INTO steps VALUES (?, ?)"); + step.run( + 0, + new Uint8Array(protoBytes(9, [...protoNumber(2, 100), ...protoText(11, "response")])), + ); + step.run( + 1, + new Uint8Array(protoBytes(9, [...protoNumber(3, 40), ...protoText(12, "provider")])), + ); + db.prepare("INSERT INTO gen_metadata VALUES (?, ?)").run( + 0, + new Uint8Array( + protoBytes(1, [ + ...protoText(19, "Gemini 3 Pro"), + ...protoBytes(4, [ + ...protoNumber(2, 50), + ...protoNumber(5, 20), + ...protoText(11, "response"), + ...protoText(12, "provider"), + ]), + ]), + ), + ); + } finally { + db.close(); + } + const result = await readAntigravityUsage(dir, 0); + assert.deepStrictEqual(result.errors, []); + const records = result.files.flatMap((file) => file.records); + assert.strictEqual(records.length, 1); + assert.deepStrictEqual(records[0]?.totals, { + uncachedInputTokens: 100, + cachedInputTokens: 20, + cacheCreationTokens: 0, + outputTokens: 40, + reasoningTokens: 0, + }); + }); + + it("merges Antigravity provider and message aliases across configured roots while keeping original ownership", async () => { + const roots = [NodePath.join(dir, "first"), NodePath.join(dir, "second")]; + for (const [index, root] of roots.entries()) { + await NodeFSP.mkdir(root); + const db = new NodeSqlite.DatabaseSync(NodePath.join(root, `session-${index}.db`)); + try { + db.exec("CREATE TABLE steps (idx INTEGER, metadata BLOB)"); + for (const identity of [7, 12]) { + const usage = [ + ...protoNumber(1, 246), + ...protoNumber(2, index === 0 ? 100 : 150), + ...protoText(11, `response-${index}-${identity}`), + ...protoText(identity, `shared-${identity}`), + ]; + db.prepare("INSERT INTO steps VALUES (?, ?)").run( + identity, + new Uint8Array(protoBytes(9, usage)), + ); + } + } finally { + db.close(); + } + } + const result = await readAntigravityUsage(roots, 0); + assert.deepStrictEqual(result.errors, []); + assert.strictEqual(result.files.length, 2); + assert.strictEqual(result.files[0]?.root, roots[0]); + assert.strictEqual(result.files[0]?.records.length, 2); + assert.strictEqual(result.files[1]?.records.length, 0); + assert.deepStrictEqual( + result.files[0]?.records.map((record) => record.totals.uncachedInputTokens), + [150, 150], + ); + assert.isTrue(result.files[0]?.records.every((record) => record.sessionId === "session-0")); + }); + + it("upgrades Antigravity fallback timestamps before applying the date window", async () => { + for (const fallback of ["mtime", "trajectory"]) { + const path = NodePath.join(dir, `${fallback}.db`); + const db = new NodeSqlite.DatabaseSync(path); + try { + db.exec( + "CREATE TABLE gen_metadata (idx INTEGER, data BLOB); CREATE TABLE steps (idx INTEGER, metadata BLOB)", + ); + if (fallback === "trajectory") { + db.exec("CREATE TABLE trajectory_metadata_blob (data BLOB)"); + db.prepare("INSERT INTO trajectory_metadata_blob VALUES (?)").run( + new Uint8Array(protoBytes(2, protoNumber(1, 1780000200))), + ); + } + for (const [index, seconds] of [1780000000, 1780000200].entries()) { + const usage = [...protoNumber(2, 10), ...protoText(11, `${fallback}-${index}`)]; + db.prepare("INSERT INTO steps VALUES (?, ?)").run( + index, + new Uint8Array(protoBytes(9, usage)), + ); + db.prepare("INSERT INTO gen_metadata VALUES (?, ?)").run( + index, + new Uint8Array( + protoBytes(1, [ + ...protoBytes(4, usage), + ...protoBytes(9, protoBytes(4, protoNumber(1, seconds))), + ]), + ), + ); + } + } finally { + db.close(); + } + await NodeFSP.utimes(path, 1780000000, 1780000000); + } + const result = await readAntigravityUsage(dir, 1780000100000); + assert.deepStrictEqual(result.errors, []); + const records = result.files.flatMap((file) => file.records); + assert.strictEqual(records.length, 2); + assert.deepStrictEqual( + records.map((record) => record.timestampMs), + [1780000200000, 1780000200000], + ); + }); + + const antigravityGeneration = (responseId: string) => + new Uint8Array( + protoBytes(1, [ + ...protoBytes(4, [ + ...protoNumber(2, 100), + ...protoNumber(3, 40), + ...protoText(11, responseId), + ]), + ...protoText(19, "Gemini 3 Pro"), + ...protoBytes(9, protoBytes(4, protoNumber(1, 1780000000))), + ]), + ); + + it("reuses an unchanged Antigravity database instead of decoding it again", async () => { + const path = NodePath.join(dir, "session-1.db"); + const db = new NodeSqlite.DatabaseSync(path); + try { + db.exec("CREATE TABLE gen_metadata (idx INTEGER, data BLOB)"); + db.prepare("INSERT INTO gen_metadata VALUES (?, ?)").run(0, antigravityGeneration("r-1")); + const cache = makeAntigravityUsageCache(); + assert.deepStrictEqual((await readAntigravityUsage(dir, 0, cache)).errors, []); + + // An exclusive lock makes a fresh read fail without touching the file, so + // only a cache hit can still return the earlier records. + db.exec("BEGIN EXCLUSIVE"); + assert.strictEqual((await readAntigravityUsage(dir, 0)).errors.length, 1); + const cached = await readAntigravityUsage(dir, 0, cache); + assert.deepStrictEqual(cached.errors, []); + assert.deepStrictEqual( + cached.files.flatMap((file) => file.records).map((record) => record.dedupeKey), + ["antigravity:11:r-1"], + ); + db.exec("ROLLBACK"); + } finally { + db.close(); + } + }); + + it("rereads an Antigravity database rewritten with its size and mtime restored", async () => { + const path = NodePath.join(dir, "session-1.db"); + const db = new NodeSqlite.DatabaseSync(path); + try { + db.exec("CREATE TABLE gen_metadata (idx INTEGER, data BLOB)"); + db.prepare("INSERT INTO gen_metadata VALUES (?, ?)").run(0, antigravityGeneration("r-1")); + } finally { + db.close(); + } + await NodeFSP.utimes(path, 1780000000, 1780000000); + const cache = makeAntigravityUsageCache(); + assert.deepStrictEqual((await readAntigravityUsage(dir, 0, cache)).errors, []); + + // ctime has the kernel's timestamp granularity, which can be a few + // milliseconds, so repeat the forged rewrite until it lands on a later tick + // than the cached read, as any real rewrite does. + const cached = await NodeFSP.stat(path); + do { + await NodeFSP.writeFile(path, Buffer.alloc(cached.size)); + await NodeFSP.utimes(path, 1780000000, 1780000000); + } while ((await NodeFSP.stat(path)).ctimeMs === cached.ctimeMs); + const restored = await NodeFSP.stat(path); + assert.strictEqual(restored.size, cached.size); + assert.strictEqual(restored.mtimeMs, cached.mtimeMs); + const next = await readAntigravityUsage(dir, 0, cache); + assert.deepStrictEqual(next.errors, [path]); + assert.deepStrictEqual( + next.files.flatMap((file) => file.records), + [], + ); + }); + + it("rereads an Antigravity database when only its WAL changed", async () => { + const path = NodePath.join(dir, "session-1.db"); + const db = new NodeSqlite.DatabaseSync(path); + try { + db.exec( + "PRAGMA journal_mode = WAL; PRAGMA wal_autocheckpoint = 0; CREATE TABLE gen_metadata (idx INTEGER, data BLOB)", + ); + const insert = db.prepare("INSERT INTO gen_metadata VALUES (?, ?)"); + insert.run(0, antigravityGeneration("r-1")); + const cache = makeAntigravityUsageCache(); + const first = await readAntigravityUsage(dir, 0, cache); + assert.strictEqual(first.files.flatMap((file) => file.records).length, 1); + + const before = await NodeFSP.stat(path); + insert.run(1, antigravityGeneration("r-2")); + const after = await NodeFSP.stat(path); + assert.strictEqual(after.size, before.size); + assert.strictEqual(after.mtimeMs, before.mtimeMs); + + const next = await readAntigravityUsage(dir, 0, cache); + assert.deepStrictEqual( + next.files.flatMap((file) => file.records).map((record) => record.dedupeKey), + ["antigravity:11:r-1", "antigravity:11:r-2"], + ); + } finally { + db.close(); + } + }); + + it("reads Antigravity step-only stores and reports malformed databases", async () => { + const db = new NodeSqlite.DatabaseSync(NodePath.join(dir, "steps.db")); + try { + db.exec("CREATE TABLE steps (idx INTEGER, metadata BLOB)"); + const usage = [...protoNumber(1, 246), ...protoNumber(2, 10), ...protoNumber(3, 5)]; + db.prepare("INSERT INTO steps VALUES (?, ?)").run( + 0, + new Uint8Array([...protoBytes(9, usage), ...protoBytes(8, protoNumber(1, 1780000000))]), + ); + } finally { + db.close(); + } + await NodeFSP.writeFile(NodePath.join(dir, "broken.db"), "not a sqlite database"); + const result = await readAntigravityUsage(dir, 0); + assert.strictEqual(result.errors.length, 1); + assert.strictEqual(result.files.flatMap((file) => file.records)[0]?.model, "gemini-2.5-pro"); + assert.strictEqual(result.files.flatMap((file) => file.records)[0]?.totals.outputTokens, 5); + }); + + it("ignores large values in unused Antigravity protobuf fields", async () => { + const db = new NodeSqlite.DatabaseSync(NodePath.join(dir, "large-varint.db")); + try { + db.exec("CREATE TABLE steps (idx INTEGER, metadata BLOB)"); + const unusedField = [...protoNumber(99, 0).slice(0, -1), ...Array(9).fill(0xff), 0x01]; + const usage = [...protoNumber(1, 246), ...protoNumber(2, 10), ...unusedField]; + db.prepare("INSERT INTO steps VALUES (?, ?)").run(0, new Uint8Array(protoBytes(9, usage))); + } finally { + db.close(); + } + const result = await readAntigravityUsage(dir, 0); + assert.deepStrictEqual(result.errors, []); + assert.strictEqual( + result.files.flatMap((file) => file.records)[0]?.totals.uncachedInputTokens, + 10, + ); + }); +}); diff --git a/apps/server/src/provider/Drivers/antigravityUsageReader.ts b/apps/server/src/provider/Drivers/antigravityUsageReader.ts new file mode 100644 index 0000000000000000000000000000000000000000..9e91b6eacbb6fb94d7942818cbfed071a8e2d7cb --- /dev/null +++ b/apps/server/src/provider/Drivers/antigravityUsageReader.ts @@ -0,0 +1,406 @@ +// node:sqlite reads live conversation databases while Node fs discovers them. +// @effect-diagnostics nodeBuiltinImport:off +import * as NodeFSP from "node:fs/promises"; +import * as NodePath from "node:path"; +import * as NodeSqlite from "node:sqlite"; +import * as NodeTimersPromises from "node:timers/promises"; + +import type { UsageRecord } from "@t3tools/provider-core/server/usage"; + +type FieldValue = number | bigint | Uint8Array; +type Fields = Map; + +/** Antigravity stores usage metadata as protobuf, independently of conversation text. */ +function fields(bytes: Uint8Array): Fields { + let offset = 0; + const result: Fields = new Map(); + const varint = () => { + let value = 0n; + for (let shift = 0n; shift < 70n; shift += 7n) { + const byte = bytes[offset++]; + if (byte === undefined || (shift === 63n && byte > 1)) { + throw new Error("Invalid Antigravity protobuf varint"); + } + value |= BigInt(byte & 127) << shift; + if (byte < 128) { + return value > BigInt(Number.MAX_SAFE_INTEGER) ? value : Number(value); + } + } + throw new Error("Invalid Antigravity protobuf varint"); + }; + while (offset < bytes.length) { + const tag = varint(); + if (typeof tag !== "number") throw new Error("Invalid protobuf field"); + const number = Math.floor(tag / 8); + const wire = tag % 8; + if (number === 0) throw new Error("Invalid protobuf field"); + let value: FieldValue; + if (wire === 0) { + value = varint(); + } else if (wire === 1 || wire === 5 || wire === 2) { + const length = wire === 2 ? varint() : wire === 1 ? 8 : 4; + if (typeof length !== "number") throw new Error("Invalid protobuf field length"); + if (length > bytes.length - offset) throw new Error("Truncated protobuf field"); + value = bytes.subarray(offset, offset + length); + offset += length; + if (wire !== 2) continue; + } else { + throw new Error("Unsupported protobuf wire type"); + } + const entries = result.get(number) ?? []; + entries.push(value); + result.set(number, entries); + } + return result; +} + +const numberAt = (value: Fields, key: number) => { + const entry = value.get(key)?.[0]; + return typeof entry === "number" ? entry : 0; +}; +const bytesAt = (value: Fields, key: number) => { + const entry = value.get(key)?.[0]; + return entry instanceof Uint8Array ? entry : undefined; +}; +const nested = (value: Fields, key: number) => { + const bytes = bytesAt(value, key); + return bytes === undefined ? new Map() : fields(bytes); +}; +const textAt = (value: Fields, key: number) => { + const bytes = bytesAt(value, key); + return bytes === undefined ? "" : new TextDecoder("utf-8", { fatal: true }).decode(bytes).trim(); +}; +const timestamp = (value: Fields) => { + const seconds = numberAt(value, 1); + return seconds > 0 ? seconds * 1000 + Math.floor(numberAt(value, 2) / 1_000_000) : null; +}; + +const MODEL_IDS: Record = { + 246: "gemini-2.5-pro", + 312: "gemini-2.5-flash", + 313: "gemini-2.5-flash-thinking", + 329: "gemini-2.5-flash-thinking", + 330: "gemini-2.5-flash-lite", + 281: "claude-sonnet-4", + 282: "claude-sonnet-4", + 290: "claude-opus-4", + 291: "claude-opus-4", + 333: "claude-sonnet-4-5", + 334: "claude-sonnet-4-5", + 340: "claude-haiku-4-5", + 341: "claude-haiku-4-5", + 1026: "claude-opus-4-6", + 1035: "claude-sonnet-4-6", + 1016: "gemini-3.1-pro", + 1036: "gemini-3.1-pro", + 1037: "gemini-3.1-pro", + 1018: "gemini-3-flash-preview", + 1084: "gemini-3-flash-preview", + 1047: "gemini-3-flash-preview", +}; + +function modelName(name: string, id: number): string { + if (name) { + const normalized = name + .toLowerCase() + .replace(/\s*\([^)]*\)\s*$/, "") + .replaceAll(" ", "-"); + if (normalized.startsWith("claude-")) { + return normalized + .replace(/^claude-(4(?:\.\d+)?)-(sonnet|opus|haiku)/, "claude-$2-$1") + .replaceAll(".", "-"); + } + return normalized; + } + return MODEL_IDS[id] ?? (id > 0 ? `antigravity-model-${id}` : ""); +} + +interface Metadata { + model: string; + timestampMs: number | null; + usages: Fields[]; +} + +function metadata(bytes: Uint8Array, step: boolean): Metadata { + const root = fields(bytes); + if (!step && bytesAt(root, 1) === undefined) { + throw new Error("Missing Antigravity generation metadata"); + } + const data = step ? root : nested(root, 1); + const model = step ? nested(data, 24) : data; + const usage = bytesAt(data, step ? 9 : 4); + const usages = usage === undefined ? [] : [fields(usage)]; + for (const retry of data.get(step ? 28 : 17) ?? []) { + if (!(retry instanceof Uint8Array)) throw new Error("Invalid retry metadata"); + const retryUsage = bytesAt(fields(retry), 2); + if (retryUsage !== undefined) usages.push(fields(retryUsage)); + } + return { + model: modelName( + textAt(model, step ? 12 : 19) || textAt(model, step ? 8 : 21), + numberAt(model, step ? 1 : 3), + ), + timestampMs: step + ? (timestamp(nested(data, 8)) ?? timestamp(nested(data, 1))) + : timestamp(nested(nested(data, 9), 4)), + usages, + }; +} + +function blob(value: unknown): Uint8Array { + if (!(value instanceof Uint8Array)) throw new Error("Invalid Antigravity metadata blob"); + return value; +} + +interface UsageCandidate { + record: UsageRecord; + keys: readonly string[]; + timestampQuality: number; +} + +async function readDatabase(path: string, fallbackTimestamp: number): Promise { + const db = new NodeSqlite.DatabaseSync(path, { readOnly: true }); + try { + db.exec("PRAGMA busy_timeout = 100; BEGIN"); + const tables = new Set( + db + .prepare("SELECT name FROM sqlite_master WHERE type = 'table'") + .all() + .map((row) => row.name), + ); + if (!tables.has("gen_metadata") && !tables.has("steps")) { + throw new Error("Missing Antigravity usage tables"); + } + const readMetadata = async (query: string, column: string, step: boolean) => { + const entries: Array<{ idx: number; entry: Metadata }> = []; + for (const row of db.prepare(query).iterate()) { + if (typeof row.idx !== "number") throw new Error("Invalid Antigravity metadata index"); + entries.push({ idx: row.idx, entry: metadata(blob(row[column]), step) }); + if (entries.length % 256 === 0) await NodeTimersPromises.setImmediate(); + } + return entries; + }; + const generations = tables.has("gen_metadata") + ? await readMetadata("SELECT idx, data FROM gen_metadata ORDER BY idx", "data", false) + : []; + let trajectoryTimestamp: number | null = null; + if (tables.has("trajectory_metadata_blob")) { + for (const row of db.prepare("SELECT data FROM trajectory_metadata_blob").iterate()) { + trajectoryTimestamp ??= timestamp(nested(fields(blob(row.data)), 2)); + } + } + const steps = tables.has("steps") + ? await readMetadata( + "SELECT idx, metadata FROM steps WHERE metadata IS NOT NULL ORDER BY idx", + "metadata", + true, + ) + : []; + const sessionId = NodePath.basename(path, ".db"); + const records: UsageCandidate[] = []; + const generationModels = new Map(generations.map(({ idx, entry }) => [idx, entry.model])); + for (const [source, entries] of [ + ["step", steps], + ["generation", generations], + ] as const) { + for (const [index, { idx, entry }] of entries.entries()) { + for (const [usageIndex, usage] of entry.usages.entries()) { + const outputTokens = Math.max( + numberAt(usage, 3), + numberAt(usage, 9) + numberAt(usage, 10), + ); + const totals = { + uncachedInputTokens: numberAt(usage, 2), + cachedInputTokens: numberAt(usage, 5), + cacheCreationTokens: numberAt(usage, 4), + outputTokens, + reasoningTokens: Math.min(outputTokens, numberAt(usage, 9)), + }; + if ( + totals.uncachedInputTokens + + totals.cachedInputTokens + + totals.cacheCreationTokens + + outputTokens === + 0 + ) + continue; + const keys = ([11, 12, 7] as const).flatMap((key) => { + const id = textAt(usage, key); + return id ? [`antigravity:${key}:${id}`] : []; + }); + const record: UsageRecord = { + provider: "antigravity", + sessionId, + timestampMs: entry.timestampMs ?? trajectoryTimestamp ?? fallbackTimestamp, + model: + MODEL_IDS[numberAt(usage, 1)] || + entry.model || + (source === "step" ? generationModels.get(idx) : "") || + modelName("", numberAt(usage, 1)) || + "antigravity-unknown", + totals, + reportedCostUsd: null, + speed: "standard", + dedupeKey: keys[0] ?? `antigravity:${sessionId}:${source}:${index}:${usageIndex}`, + }; + records.push({ + record, + keys, + timestampQuality: entry.timestampMs !== null ? 2 : trajectoryTimestamp !== null ? 1 : 0, + }); + } + } + } + return records; + } finally { + db.close(); + } +} + +interface CachedDatabase { + readonly fingerprint: string; + readonly candidates: readonly UsageCandidate[]; +} + +/** + * Parsed databases keyed by canonical path. An entry is reused only while both + * the database and its `-wal` sidecar keep the same size, mtime and ctime: new + * rows land in the WAL without touching the main file until a checkpoint, and + * ctime moves on any content write even when mtime is restored. + */ +export const makeAntigravityUsageCache = () => new Map(); + +/** + * Reads and merges aliases across every configured Antigravity store before date + * filtering. With a cache, databases unchanged since the previous read are not + * decoded again. + */ +export async function readAntigravityUsage( + conversationsDirectories: string | readonly string[], + sinceMs: number, + cache?: Map, +) { + const roots = + typeof conversationsDirectories === "string" + ? [conversationsDirectories] + : conversationsDirectories; + const files: Array<{ root: string; path: string; records: UsageRecord[] }> = []; + const errors: string[] = []; + const identities = new Map(); + const groups: Array< + UsageCandidate & { parent: number; size: number; owner: number; fileIndex: number } + > = []; + const find = (index: number): number => { + let root = index; + while (groups[root]!.parent !== root) root = groups[root]!.parent; + while (index !== root) { + const parent = groups[index]!.parent; + groups[index]!.parent = root; + index = parent; + } + return root; + }; + const merge = (left: number, right: number): number => { + let a = find(left); + let b = find(right); + if (a === b) return a; + if (groups[a]!.size < groups[b]!.size) [a, b] = [b, a]; + const target = groups[a]!; + const source = groups[b]!; + const first = target.owner < source.owner ? target : source; + const bestTime = + source.timestampQuality > target.timestampQuality || + (source.timestampQuality === target.timestampQuality && + source.record.timestampMs < target.record.timestampMs) + ? source + : target; + const x = target.record.totals; + const y = source.record.totals; + target.record = { + ...first.record, + model: + first.record.model === "antigravity-unknown" + ? first === target + ? source.record.model + : target.record.model + : first.record.model, + timestampMs: bestTime.record.timestampMs, + totals: { + uncachedInputTokens: Math.max(x.uncachedInputTokens, y.uncachedInputTokens), + cachedInputTokens: Math.max(x.cachedInputTokens, y.cachedInputTokens), + cacheCreationTokens: Math.max(x.cacheCreationTokens, y.cacheCreationTokens), + outputTokens: Math.max(x.outputTokens, y.outputTokens), + reasoningTokens: Math.max(x.reasoningTokens, y.reasoningTokens), + }, + }; + target.timestampQuality = bestTime.timestampQuality; + target.owner = first.owner; + target.fileIndex = first.fileIndex; + target.size += source.size; + source.parent = a; + return a; + }; + const append = (candidate: UsageCandidate, fileIndex: number) => { + const index = groups.length; + groups.push({ ...candidate, parent: index, size: 1, owner: index, fileIndex }); + for (const key of candidate.keys) { + const existing = identities.get(key); + if (existing !== undefined) merge(index, existing); + identities.set(key, index); + } + }; + const visited = new Set(); + const walk = async (directory: string, root: string): Promise => { + let entries; + try { + entries = await NodeFSP.readdir(directory, { withFileTypes: true }); + } catch (error) { + if ((error as NodeJS.ErrnoException).code !== "ENOENT") errors.push(directory); + return; + } + entries.sort((a, b) => a.name.localeCompare(b.name)); + for (const entry of entries) { + const path = NodePath.join(directory, entry.name); + if (entry.isDirectory()) { + await walk(path, root); + } else if (entry.isFile() && entry.name.endsWith(".db")) { + try { + const canonical = await NodeFSP.realpath(path); + if (visited.has(canonical)) continue; + visited.add(canonical); + const stat = await NodeFSP.stat(path); + const wal = await NodeFSP.stat(`${path}-wal`).catch(() => null); + const fingerprint = [stat, wal] + .map((file) => (file ? `${file.size}:${file.mtimeMs}:${file.ctimeMs}` : "-")) + .join("/"); + const cached = cache?.get(canonical); + let candidates: readonly UsageCandidate[]; + if (cached?.fingerprint === fingerprint) { + candidates = cached.candidates; + } else { + candidates = await readDatabase(path, stat.mtimeMs); + cache?.set(canonical, { fingerprint, candidates }); + } + const fileIndex = files.length; + files.push({ root, path, records: [] }); + for (const [index, candidate] of candidates.entries()) { + append(candidate, fileIndex); + if (index % 256 === 255) await NodeTimersPromises.setImmediate(); + } + } catch { + errors.push(path); + } + } + } + }; + for (const root of roots) await walk(root, root); + if (cache !== undefined) { + for (const key of cache.keys()) if (!visited.has(key)) cache.delete(key); + } + for (const [index, group] of groups.entries()) { + if (group.parent === index && group.record.timestampMs >= sinceMs) { + files[group.fileIndex]!.records.push(group.record); + } + } + return { files, errors }; +} diff --git a/apps/server/src/provider/builtInDrivers.ts b/apps/server/src/provider/builtInDrivers.ts index 9279ef326a259d6e03a3c22675d268b999a04ace..532bd602d1453cd3d1906fd9f3ba4ad2e7ce2547 100644 --- a/apps/server/src/provider/builtInDrivers.ts +++ b/apps/server/src/provider/builtInDrivers.ts @@ -76,6 +76,7 @@ export type BuiltInUsageReadersEnv = | ProviderUsageReaderEnv | ProviderUsageReaderEnv | ProviderUsageReaderEnv + | ProviderUsageReaderEnv | ProviderUsageReaderEnv; /** @@ -85,4 +86,4 @@ export type BuiltInUsageReadersEnv = */ export const BUILT_IN_USAGE_DRIVERS: ReadonlyArray< AnyProviderDriver -> = [ClaudeDriver, CodexDriver, GrokDriver, OpenCodeDriver, CursorDriver]; +> = [ClaudeDriver, CodexDriver, GrokDriver, OpenCodeDriver, AntigravityDriver, CursorDriver]; diff --git a/apps/server/src/server.ts b/apps/server/src/server.ts index 5659b130168d560fca0272876fd03999b3f1238b..80c5a73ac4c185511c156a7d25d22ce5491ec75b 100644 --- a/apps/server/src/server.ts +++ b/apps/server/src/server.ts @@ -154,6 +154,7 @@ import * as NativeTelemetryClient from "./resourceTelemetry/NativeTelemetryClien import * as ResourceAttribution from "./resourceTelemetry/ResourceAttribution.ts"; import * as ResourceMonitorBinary from "./resourceTelemetry/ResourceMonitorBinary.ts"; import * as ResourceTelemetry from "./resourceTelemetry/ResourceTelemetry.ts"; +import * as AntigravityUsage from "./provider/Drivers/AntigravityUsage.ts"; import * as CursorAccountReader from "@t3tools/provider-cursor/server/CursorAccountReader"; import * as CursorKeychain from "@t3tools/provider-cursor/server/CursorKeychain"; import * as CursorUsageAccounts from "@t3tools/provider-cursor/server/CursorUsageAccounts"; @@ -255,10 +256,12 @@ const layerBackground = BackgroundPolicy.layer.pipe( const layerUsage = UsageService.layer.pipe( Layer.provide( - CursorUsageAccounts.layer.pipe( - Layer.provide(CursorAccountReader.layer.pipe(Layer.provide(CursorKeychain.layer))), - Layer.provide(ProviderHostLive.layer.pipe(Layer.provide(ServerSecretStore.layer))), - ), + Layer.mergeAll( + AntigravityUsage.layer, + CursorUsageAccounts.layer.pipe( + Layer.provide(CursorAccountReader.layer.pipe(Layer.provide(CursorKeychain.layer))), + ), + ).pipe(Layer.provide(ProviderHostLive.layer.pipe(Layer.provide(ServerSecretStore.layer)))), ), Layer.provide(layerServerSettings), ); diff --git a/apps/server/src/usage/UsageService.test.ts b/apps/server/src/usage/UsageService.test.ts index 83bace94138afbbcfe4f06ecb4d48322e2b0c4dd..68d2e7c2bbf6d7938cb2bbf74f288be9c1ab4210 100644 --- a/apps/server/src/usage/UsageService.test.ts +++ b/apps/server/src/usage/UsageService.test.ts @@ -38,6 +38,7 @@ import * as CursorKeychain from "@t3tools/provider-cursor/server/CursorKeychain" import * as CursorUsageAccounts from "@t3tools/provider-cursor/server/CursorUsageAccounts"; import * as ServerSecretStore from "../auth/ServerSecretStore.ts"; import * as BackgroundPolicy from "../background/BackgroundPolicy.ts"; +import * as AntigravityUsage from "../provider/Drivers/AntigravityUsage.ts"; import * as ProviderHostLive from "../provider/ProviderHostLive.ts"; import type { UsageRecord } from "@t3tools/provider-core/server/usage"; import * as UsageService from "./UsageService.ts"; @@ -127,6 +128,7 @@ const layerService = (input: { layerCursorUsageAccounts( CursorAccountReader.layer.pipe(Layer.provide(CursorKeychain.layer)), ).pipe( + Layer.provideMerge(AntigravityUsage.layer), Layer.provideMerge(ProviderHostLive.layer), Layer.provideMerge(Layer.mock(BackgroundPolicy.BackgroundPolicy)({})), Layer.provideMerge(Layer.mock(ServerSecretStore.ServerSecretStore)({})), diff --git a/apps/server/src/usage/UsageService.ts b/apps/server/src/usage/UsageService.ts index 2f48fc5abcab6721f2459c41e51e81502d444f8a..ba65ba2887bbcab0328c287471f5164977e8397a 100644 --- a/apps/server/src/usage/UsageService.ts +++ b/apps/server/src/usage/UsageService.ts @@ -1,16 +1,15 @@ /** - * UsageService - scans provider transcripts and returns priced usage buckets. + * UsageService - reads every driver's usage history and returns priced usage + * buckets. * - * The scan reads native session files and databases, including work driven - * outside T3 Code. Cursor's local records provide only partial coverage. + * Each built-in driver with a `usage` reader contributes its sources, including + * work driven outside T3 Code. A `transcripts` reader names JSONL directories + * that this service streams itself; a `scan` reader reads its own sources. * * JSONL transcripts are append-only, so parsed records are memoised per file by * `(size, mtime)`. A cold 30-day scan of ~1.4 GB lands around 2-3 seconds; warm * scans only reparse files that changed, and a file that merely grew resumes * from its cached parse position so only the appended bytes are read. - * OpenCode's SQLite reader queries the live database each scan so WAL writes - * remain visible. Antigravity databases are memoised in memory while the - * database and its WAL keep the same `(size, mtime, ctime)`. * * A scan reader with a slow source (an account API) may answer from its own * cache and mark itself `refreshing` while a background refresh runs; @@ -36,7 +35,6 @@ import { HostProcessEnvironment } from "@t3tools/shared/hostProcess"; import * as Cause from "effect/Cause"; import * as Clock from "effect/Clock"; import * as Context from "effect/Context"; -import * as Crypto from "effect/Crypto"; import * as DateTime from "effect/DateTime"; import * as Deferred from "effect/Deferred"; import * as Effect from "effect/Effect"; @@ -52,9 +50,7 @@ import { HttpClient, HttpClientResponse } from "effect/http"; import { writeFileStringAtomically } from "@t3tools/shared/atomicWrite"; import * as ServerConfig from "../config.ts"; -import { expandHomePath } from "@t3tools/provider-core/server/pathExpansion"; import * as ServerSettings from "../serverSettings.ts"; -import { resolveAntigravityInstanceDirectories } from "../provider/antigravityAuthSupport.ts"; import { BUILT_IN_USAGE_DRIVERS, type BuiltInUsageReadersEnv } from "../provider/builtInDrivers.ts"; import type { ProviderDriver } from "@t3tools/provider-core/server/driver"; import { mergeProviderInstanceEnvironment } from "@t3tools/provider-core/server/instanceEnvironment"; @@ -63,7 +59,6 @@ import type { TranscriptUsageFormat, UsageRecord, } from "@t3tools/provider-core/server/usage"; -import { makeAntigravityUsageCache, readAntigravityUsage } from "./antigravityUsageReader.ts"; import { resolveModelAliases, UsageAggregator } from "./usageAggregation.ts"; import { createOverrideRateTable, parseRateTable, type RateTable } from "./usagePricing.ts"; import { @@ -227,7 +222,6 @@ const layerTest = Layer.succeed( ); export const make = Effect.gen(function* () { - const crypto = yield* Crypto.Crypto; const fileSystem = yield* FileSystem.FileSystem; const path = yield* Path.Path; const config = yield* ServerConfig.ServerConfig; @@ -238,7 +232,6 @@ export const make = Effect.gen(function* () { const readerContext = yield* Effect.context(); const fileCache: ScanCache = new Map(); - const antigravityCache = makeAntigravityUsageCache(); const sourceCache = new Map(); let cacheDirty = false; const isWithinDirectory = (filePath: string, dir: string) => { @@ -639,22 +632,6 @@ export const make = Effect.gen(function* () { Effect.provideContext(readerContext), ); - const home = NodeOS.homedir(); - const envRoots = Effect.fnUntraced(function* (key: string, defaults: readonly string[]) { - const roots = hostEnvironment[key] - ?.split(",") - .map((value) => value.trim()) - .filter(Boolean); - const canonical = new Set(); - for (const root of roots?.length ? roots : defaults) { - const resolved = path.resolve(expandHomePath(root)); - canonical.add( - yield* fileSystem.realPath(resolved).pipe(Effect.orElseSucceed(() => resolved)), - ); - } - return [...canonical]; - }); - const scans = Effect.forEach( scanReaders, ({ driver, reader }) => @@ -686,83 +663,14 @@ export const make = Effect.gen(function* () { { concurrency: "unbounded" }, ); - const antigravity = Effect.gen(function* () { - const antigravityRoots = yield* envRoots("ANTIGRAVITY_DATA_DIR", [ - ...["antigravity", "antigravity-cli", "antigravity-ide", "antigravity-backup"].map((name) => - path.join(home, ".gemini", name), - ), - path.join(home, ".config", "antigravity"), - ]); - for (const [instanceId, instance] of Object.entries(settings.providerInstances)) { - if (instance.driver === "antigravity") { - const directories = yield* resolveAntigravityInstanceDirectories( - config.stateDir, - ProviderInstanceId.make(instanceId), - ).pipe( - Effect.provideService(Crypto.Crypto, crypto), - Effect.provideService(Path.Path, path), - Effect.mapError( - (cause) => - new UsageReadError({ - reason: "scanFailed", - detail: "Antigravity profile directory could not be resolved.", - cause, - }), - ), - ); - antigravityRoots.push(path.join(directories.profile, "antigravity-acp")); - } - } - const antigravityDirs = new Set(); - for (const root of antigravityRoots) { - const resolvedRoot = yield* fileSystem - .realPath(root) - .pipe(Effect.orElseSucceed(() => root)); - const nested = path.join(resolvedRoot, "conversations"); - const dir = (yield* fileSystem - .exists(nested) - .pipe(Effect.catchCause(() => Effect.succeed(false)))) - ? nested - : resolvedRoot; - antigravityDirs.add(yield* fileSystem.realPath(dir).pipe(Effect.orElseSucceed(() => dir))); - } - const result = yield* Effect.promise(() => - readAntigravityUsage([...antigravityDirs], windowStartMs, antigravityCache), - ); - const scanned: ScannedDir[] = []; - for (const dir of antigravityDirs) { - const exists = yield* fileSystem - .exists(dir) - .pipe(Effect.catchCause(() => Effect.succeed(false))); - const failed = result.errors.some( - (error) => error === dir || error.startsWith(`${dir}${path.sep}`), - ); - scanned.push({ - provider: "antigravity", - dir, - volumeId: yield* Effect.promise(() => readDirectoryVolumeId(dir)), - files: !exists && !failed ? null : result.files.filter((file) => file.root === dir), - status: failed ? "partial" : "ok", - ...(failed ? { message: "Some Antigravity history could not be read." } : {}), - }); - } - return scanned; - }); - // Independent sources scan together. Transcript directories go one at a // time, so open files stay at `TRANSCRIPT_READ_CONCURRENCY`. The result // keeps this order, since aggregation keeps the first copy of a duplicate. - const [transcripts, scanDirs, antigravityDirs] = yield* Effect.all( - [Effect.forEach(dirs, (dir) => scanTranscriptDir(dir, windowStartMs)), scans, antigravity], + const [transcripts, scanDirs] = yield* Effect.all( + [Effect.forEach(dirs, (dir) => scanTranscriptDir(dir, windowStartMs)), scans], { concurrency: "unbounded" }, ); - // Antigravity reads between the other scan readers until it has a reader of its own. - const scanned: readonly ScannedDir[] = [ - ...transcripts, - ...scanDirs.slice(0, -1).flat(), - ...antigravityDirs, - ...(scanDirs.at(-1) ?? []), - ]; + const scanned: readonly ScannedDir[] = [...transcripts, ...scanDirs.flat()]; return scanned; }); @@ -944,9 +852,10 @@ export const make = Effect.gen(function* () { }); /** - * In-flight scans by window and usage settings, so concurrent identical requests (the usage + * In-flight scans by window and settings, so concurrent identical requests (the usage * page open on two clients at once) share one scan instead of racing over - * the same corpus twice. + * the same corpus twice. Readers read settings of their own, so a scan is + * shared only under the same settings snapshot. */ const inflightScans = new Map>(); @@ -958,9 +867,7 @@ export const make = Effect.gen(function* () { input.resolution ?? "day", input.sinceTime ?? null, input.untilTime ?? null, - settings.usagePriceOverrides, - settings.usageModelAliases, - settings.cursorKeychainUsageEnabled, + settings, // A waiting read must never share a scan that answers with `refreshing`. input.awaitRefresh === true, ]); diff --git a/apps/server/src/usage/antigravityUsageReader.ts b/apps/server/src/usage/antigravityUsageReader.ts deleted file mode 100644 index 9e91b6eacbb6fb94d7942818cbfed071a8e2d7cb..0000000000000000000000000000000000000000 --- a/apps/server/src/usage/antigravityUsageReader.ts +++ /dev/null @@ -1,406 +0,0 @@ -// node:sqlite reads live conversation databases while Node fs discovers them. -// @effect-diagnostics nodeBuiltinImport:off -import * as NodeFSP from "node:fs/promises"; -import * as NodePath from "node:path"; -import * as NodeSqlite from "node:sqlite"; -import * as NodeTimersPromises from "node:timers/promises"; - -import type { UsageRecord } from "@t3tools/provider-core/server/usage"; - -type FieldValue = number | bigint | Uint8Array; -type Fields = Map; - -/** Antigravity stores usage metadata as protobuf, independently of conversation text. */ -function fields(bytes: Uint8Array): Fields { - let offset = 0; - const result: Fields = new Map(); - const varint = () => { - let value = 0n; - for (let shift = 0n; shift < 70n; shift += 7n) { - const byte = bytes[offset++]; - if (byte === undefined || (shift === 63n && byte > 1)) { - throw new Error("Invalid Antigravity protobuf varint"); - } - value |= BigInt(byte & 127) << shift; - if (byte < 128) { - return value > BigInt(Number.MAX_SAFE_INTEGER) ? value : Number(value); - } - } - throw new Error("Invalid Antigravity protobuf varint"); - }; - while (offset < bytes.length) { - const tag = varint(); - if (typeof tag !== "number") throw new Error("Invalid protobuf field"); - const number = Math.floor(tag / 8); - const wire = tag % 8; - if (number === 0) throw new Error("Invalid protobuf field"); - let value: FieldValue; - if (wire === 0) { - value = varint(); - } else if (wire === 1 || wire === 5 || wire === 2) { - const length = wire === 2 ? varint() : wire === 1 ? 8 : 4; - if (typeof length !== "number") throw new Error("Invalid protobuf field length"); - if (length > bytes.length - offset) throw new Error("Truncated protobuf field"); - value = bytes.subarray(offset, offset + length); - offset += length; - if (wire !== 2) continue; - } else { - throw new Error("Unsupported protobuf wire type"); - } - const entries = result.get(number) ?? []; - entries.push(value); - result.set(number, entries); - } - return result; -} - -const numberAt = (value: Fields, key: number) => { - const entry = value.get(key)?.[0]; - return typeof entry === "number" ? entry : 0; -}; -const bytesAt = (value: Fields, key: number) => { - const entry = value.get(key)?.[0]; - return entry instanceof Uint8Array ? entry : undefined; -}; -const nested = (value: Fields, key: number) => { - const bytes = bytesAt(value, key); - return bytes === undefined ? new Map() : fields(bytes); -}; -const textAt = (value: Fields, key: number) => { - const bytes = bytesAt(value, key); - return bytes === undefined ? "" : new TextDecoder("utf-8", { fatal: true }).decode(bytes).trim(); -}; -const timestamp = (value: Fields) => { - const seconds = numberAt(value, 1); - return seconds > 0 ? seconds * 1000 + Math.floor(numberAt(value, 2) / 1_000_000) : null; -}; - -const MODEL_IDS: Record = { - 246: "gemini-2.5-pro", - 312: "gemini-2.5-flash", - 313: "gemini-2.5-flash-thinking", - 329: "gemini-2.5-flash-thinking", - 330: "gemini-2.5-flash-lite", - 281: "claude-sonnet-4", - 282: "claude-sonnet-4", - 290: "claude-opus-4", - 291: "claude-opus-4", - 333: "claude-sonnet-4-5", - 334: "claude-sonnet-4-5", - 340: "claude-haiku-4-5", - 341: "claude-haiku-4-5", - 1026: "claude-opus-4-6", - 1035: "claude-sonnet-4-6", - 1016: "gemini-3.1-pro", - 1036: "gemini-3.1-pro", - 1037: "gemini-3.1-pro", - 1018: "gemini-3-flash-preview", - 1084: "gemini-3-flash-preview", - 1047: "gemini-3-flash-preview", -}; - -function modelName(name: string, id: number): string { - if (name) { - const normalized = name - .toLowerCase() - .replace(/\s*\([^)]*\)\s*$/, "") - .replaceAll(" ", "-"); - if (normalized.startsWith("claude-")) { - return normalized - .replace(/^claude-(4(?:\.\d+)?)-(sonnet|opus|haiku)/, "claude-$2-$1") - .replaceAll(".", "-"); - } - return normalized; - } - return MODEL_IDS[id] ?? (id > 0 ? `antigravity-model-${id}` : ""); -} - -interface Metadata { - model: string; - timestampMs: number | null; - usages: Fields[]; -} - -function metadata(bytes: Uint8Array, step: boolean): Metadata { - const root = fields(bytes); - if (!step && bytesAt(root, 1) === undefined) { - throw new Error("Missing Antigravity generation metadata"); - } - const data = step ? root : nested(root, 1); - const model = step ? nested(data, 24) : data; - const usage = bytesAt(data, step ? 9 : 4); - const usages = usage === undefined ? [] : [fields(usage)]; - for (const retry of data.get(step ? 28 : 17) ?? []) { - if (!(retry instanceof Uint8Array)) throw new Error("Invalid retry metadata"); - const retryUsage = bytesAt(fields(retry), 2); - if (retryUsage !== undefined) usages.push(fields(retryUsage)); - } - return { - model: modelName( - textAt(model, step ? 12 : 19) || textAt(model, step ? 8 : 21), - numberAt(model, step ? 1 : 3), - ), - timestampMs: step - ? (timestamp(nested(data, 8)) ?? timestamp(nested(data, 1))) - : timestamp(nested(nested(data, 9), 4)), - usages, - }; -} - -function blob(value: unknown): Uint8Array { - if (!(value instanceof Uint8Array)) throw new Error("Invalid Antigravity metadata blob"); - return value; -} - -interface UsageCandidate { - record: UsageRecord; - keys: readonly string[]; - timestampQuality: number; -} - -async function readDatabase(path: string, fallbackTimestamp: number): Promise { - const db = new NodeSqlite.DatabaseSync(path, { readOnly: true }); - try { - db.exec("PRAGMA busy_timeout = 100; BEGIN"); - const tables = new Set( - db - .prepare("SELECT name FROM sqlite_master WHERE type = 'table'") - .all() - .map((row) => row.name), - ); - if (!tables.has("gen_metadata") && !tables.has("steps")) { - throw new Error("Missing Antigravity usage tables"); - } - const readMetadata = async (query: string, column: string, step: boolean) => { - const entries: Array<{ idx: number; entry: Metadata }> = []; - for (const row of db.prepare(query).iterate()) { - if (typeof row.idx !== "number") throw new Error("Invalid Antigravity metadata index"); - entries.push({ idx: row.idx, entry: metadata(blob(row[column]), step) }); - if (entries.length % 256 === 0) await NodeTimersPromises.setImmediate(); - } - return entries; - }; - const generations = tables.has("gen_metadata") - ? await readMetadata("SELECT idx, data FROM gen_metadata ORDER BY idx", "data", false) - : []; - let trajectoryTimestamp: number | null = null; - if (tables.has("trajectory_metadata_blob")) { - for (const row of db.prepare("SELECT data FROM trajectory_metadata_blob").iterate()) { - trajectoryTimestamp ??= timestamp(nested(fields(blob(row.data)), 2)); - } - } - const steps = tables.has("steps") - ? await readMetadata( - "SELECT idx, metadata FROM steps WHERE metadata IS NOT NULL ORDER BY idx", - "metadata", - true, - ) - : []; - const sessionId = NodePath.basename(path, ".db"); - const records: UsageCandidate[] = []; - const generationModels = new Map(generations.map(({ idx, entry }) => [idx, entry.model])); - for (const [source, entries] of [ - ["step", steps], - ["generation", generations], - ] as const) { - for (const [index, { idx, entry }] of entries.entries()) { - for (const [usageIndex, usage] of entry.usages.entries()) { - const outputTokens = Math.max( - numberAt(usage, 3), - numberAt(usage, 9) + numberAt(usage, 10), - ); - const totals = { - uncachedInputTokens: numberAt(usage, 2), - cachedInputTokens: numberAt(usage, 5), - cacheCreationTokens: numberAt(usage, 4), - outputTokens, - reasoningTokens: Math.min(outputTokens, numberAt(usage, 9)), - }; - if ( - totals.uncachedInputTokens + - totals.cachedInputTokens + - totals.cacheCreationTokens + - outputTokens === - 0 - ) - continue; - const keys = ([11, 12, 7] as const).flatMap((key) => { - const id = textAt(usage, key); - return id ? [`antigravity:${key}:${id}`] : []; - }); - const record: UsageRecord = { - provider: "antigravity", - sessionId, - timestampMs: entry.timestampMs ?? trajectoryTimestamp ?? fallbackTimestamp, - model: - MODEL_IDS[numberAt(usage, 1)] || - entry.model || - (source === "step" ? generationModels.get(idx) : "") || - modelName("", numberAt(usage, 1)) || - "antigravity-unknown", - totals, - reportedCostUsd: null, - speed: "standard", - dedupeKey: keys[0] ?? `antigravity:${sessionId}:${source}:${index}:${usageIndex}`, - }; - records.push({ - record, - keys, - timestampQuality: entry.timestampMs !== null ? 2 : trajectoryTimestamp !== null ? 1 : 0, - }); - } - } - } - return records; - } finally { - db.close(); - } -} - -interface CachedDatabase { - readonly fingerprint: string; - readonly candidates: readonly UsageCandidate[]; -} - -/** - * Parsed databases keyed by canonical path. An entry is reused only while both - * the database and its `-wal` sidecar keep the same size, mtime and ctime: new - * rows land in the WAL without touching the main file until a checkpoint, and - * ctime moves on any content write even when mtime is restored. - */ -export const makeAntigravityUsageCache = () => new Map(); - -/** - * Reads and merges aliases across every configured Antigravity store before date - * filtering. With a cache, databases unchanged since the previous read are not - * decoded again. - */ -export async function readAntigravityUsage( - conversationsDirectories: string | readonly string[], - sinceMs: number, - cache?: Map, -) { - const roots = - typeof conversationsDirectories === "string" - ? [conversationsDirectories] - : conversationsDirectories; - const files: Array<{ root: string; path: string; records: UsageRecord[] }> = []; - const errors: string[] = []; - const identities = new Map(); - const groups: Array< - UsageCandidate & { parent: number; size: number; owner: number; fileIndex: number } - > = []; - const find = (index: number): number => { - let root = index; - while (groups[root]!.parent !== root) root = groups[root]!.parent; - while (index !== root) { - const parent = groups[index]!.parent; - groups[index]!.parent = root; - index = parent; - } - return root; - }; - const merge = (left: number, right: number): number => { - let a = find(left); - let b = find(right); - if (a === b) return a; - if (groups[a]!.size < groups[b]!.size) [a, b] = [b, a]; - const target = groups[a]!; - const source = groups[b]!; - const first = target.owner < source.owner ? target : source; - const bestTime = - source.timestampQuality > target.timestampQuality || - (source.timestampQuality === target.timestampQuality && - source.record.timestampMs < target.record.timestampMs) - ? source - : target; - const x = target.record.totals; - const y = source.record.totals; - target.record = { - ...first.record, - model: - first.record.model === "antigravity-unknown" - ? first === target - ? source.record.model - : target.record.model - : first.record.model, - timestampMs: bestTime.record.timestampMs, - totals: { - uncachedInputTokens: Math.max(x.uncachedInputTokens, y.uncachedInputTokens), - cachedInputTokens: Math.max(x.cachedInputTokens, y.cachedInputTokens), - cacheCreationTokens: Math.max(x.cacheCreationTokens, y.cacheCreationTokens), - outputTokens: Math.max(x.outputTokens, y.outputTokens), - reasoningTokens: Math.max(x.reasoningTokens, y.reasoningTokens), - }, - }; - target.timestampQuality = bestTime.timestampQuality; - target.owner = first.owner; - target.fileIndex = first.fileIndex; - target.size += source.size; - source.parent = a; - return a; - }; - const append = (candidate: UsageCandidate, fileIndex: number) => { - const index = groups.length; - groups.push({ ...candidate, parent: index, size: 1, owner: index, fileIndex }); - for (const key of candidate.keys) { - const existing = identities.get(key); - if (existing !== undefined) merge(index, existing); - identities.set(key, index); - } - }; - const visited = new Set(); - const walk = async (directory: string, root: string): Promise => { - let entries; - try { - entries = await NodeFSP.readdir(directory, { withFileTypes: true }); - } catch (error) { - if ((error as NodeJS.ErrnoException).code !== "ENOENT") errors.push(directory); - return; - } - entries.sort((a, b) => a.name.localeCompare(b.name)); - for (const entry of entries) { - const path = NodePath.join(directory, entry.name); - if (entry.isDirectory()) { - await walk(path, root); - } else if (entry.isFile() && entry.name.endsWith(".db")) { - try { - const canonical = await NodeFSP.realpath(path); - if (visited.has(canonical)) continue; - visited.add(canonical); - const stat = await NodeFSP.stat(path); - const wal = await NodeFSP.stat(`${path}-wal`).catch(() => null); - const fingerprint = [stat, wal] - .map((file) => (file ? `${file.size}:${file.mtimeMs}:${file.ctimeMs}` : "-")) - .join("/"); - const cached = cache?.get(canonical); - let candidates: readonly UsageCandidate[]; - if (cached?.fingerprint === fingerprint) { - candidates = cached.candidates; - } else { - candidates = await readDatabase(path, stat.mtimeMs); - cache?.set(canonical, { fingerprint, candidates }); - } - const fileIndex = files.length; - files.push({ root, path, records: [] }); - for (const [index, candidate] of candidates.entries()) { - append(candidate, fileIndex); - if (index % 256 === 255) await NodeTimersPromises.setImmediate(); - } - } catch { - errors.push(path); - } - } - } - }; - for (const root of roots) await walk(root, root); - if (cache !== undefined) { - for (const key of cache.keys()) if (!visited.has(key)) cache.delete(key); - } - for (const [index, group] of groups.entries()) { - if (group.parent === index && group.record.timestampMs >= sinceMs) { - files[group.fileIndex]!.records.push(group.record); - } - } - return { files, errors }; -} diff --git a/apps/server/src/usage/usageTranscriptReader.test.ts b/apps/server/src/usage/usageTranscriptReader.test.ts index 2cac91423697a6ccea4b701fdf206d91987f88da..e2c4ee1827feac35e39092c2ba942ffbeea98993 100644 --- a/apps/server/src/usage/usageTranscriptReader.test.ts +++ b/apps/server/src/usage/usageTranscriptReader.test.ts @@ -4,39 +4,14 @@ import * as NodeFSP from "node:fs/promises"; import * as NodeOS from "node:os"; import * as NodePath from "node:path"; -import * as NodeSqlite from "node:sqlite"; import { afterEach, assert, beforeEach, describe, it } from "@effect/vitest"; import { TEST_FORMATS } from "./usageTestFormats.ts"; import { readTranscriptRecords } from "./usageTranscriptReader.ts"; -import { makeAntigravityUsageCache, readAntigravityUsage } from "./antigravityUsageReader.ts"; let dir: string; -function protoNumber(field: number, value: number): number[] { - const varint = (number: number) => { - const bytes: number[] = []; - do { - const byte = number % 128; - number = Math.floor(number / 128); - bytes.push(byte + (number > 0 ? 128 : 0)); - } while (number > 0); - return bytes; - }; - return [...varint(field * 8), ...varint(value)]; -} - -function protoBytes(field: number, bytes: readonly number[]): number[] { - const encoded = protoNumber(field, bytes.length); - encoded[0] = encoded[0]! + 2; - return [...encoded, ...bytes]; -} - -function protoText(field: number, value: string): number[] { - return protoBytes(field, [...Buffer.from(value)]); -} - beforeEach(async () => { dir = await NodeFSP.mkdtemp(NodePath.join(NodeOS.tmpdir(), "usage-reader-test-")); }); @@ -249,351 +224,3 @@ describe("readTranscriptRecords resume", () => { ); }); }); - -describe("SQLite usage readers", () => { - it("deduplicates Antigravity generation and step usage while preserving retry model and token buckets", async () => { - const db = new NodeSqlite.DatabaseSync(NodePath.join(dir, "session-1.db")); - const stamp = protoNumber(1, 1780000000); - const usage = [ - ...protoNumber(2, 100), - ...protoNumber(3, 40), - ...protoNumber(4, 5), - ...protoNumber(5, 20), - ...protoNumber(9, 10), - ...protoText(11, "response-1"), - ]; - const retry = [ - ...protoNumber(1, 1026), - ...protoNumber(2, 12), - ...protoNumber(3, 3), - ...protoText(11, "retry-1"), - ]; - const generation = protoBytes(1, [ - ...protoBytes(4, usage), - ...protoText(19, "Gemini 3 Pro"), - ...protoBytes(9, protoBytes(4, stamp)), - ]); - const step = [ - ...protoBytes(9, usage), - ...protoBytes(8, stamp), - ...protoBytes(28, protoBytes(2, retry)), - ]; - try { - db.exec( - "CREATE TABLE gen_metadata (idx INTEGER, data BLOB); CREATE TABLE steps (idx INTEGER, metadata BLOB)", - ); - db.prepare("INSERT INTO gen_metadata VALUES (?, ?)").run(0, new Uint8Array(generation)); - db.prepare("INSERT INTO steps VALUES (?, ?)").run(0, new Uint8Array(step)); - } finally { - db.close(); - } - const result = await readAntigravityUsage(dir, 0); - assert.deepStrictEqual(result.errors, []); - const records = result.files.flatMap((file) => file.records); - assert.strictEqual(records.length, 2); - const main = records.find((record) => record.model === "gemini-3-pro"); - assert.isDefined(main); - assert.strictEqual(main?.timestampMs, 1780000000000); - assert.strictEqual(main?.sessionId, "session-1"); - assert.deepStrictEqual(main?.totals, { - uncachedInputTokens: 100, - cachedInputTokens: 20, - cacheCreationTokens: 5, - outputTokens: 40, - reasoningTokens: 10, - }); - assert.strictEqual( - records.find((record) => record.model === "claude-opus-4-6")?.totals.uncachedInputTokens, - 12, - ); - assert.deepStrictEqual( - (await readAntigravityUsage(dir, 1780000000001)).files.flatMap((file) => file.records), - [], - ); - }); - - it("uses the matching Antigravity generation model for each model-less step", async () => { - const db = new NodeSqlite.DatabaseSync(NodePath.join(dir, "model-switch.db")); - try { - db.exec( - "CREATE TABLE gen_metadata (idx INTEGER, data BLOB); CREATE TABLE steps (idx INTEGER, metadata BLOB)", - ); - const generation = db.prepare("INSERT INTO gen_metadata VALUES (?, ?)"); - const step = db.prepare("INSERT INTO steps VALUES (?, ?)"); - for (const [idx, name] of ["Gemini 3 Pro", "Claude Opus 4.6"].entries()) { - generation.run(idx, new Uint8Array(protoBytes(1, protoText(19, name)))); - step.run(idx, new Uint8Array(protoBytes(9, protoNumber(2, 10 + idx)))); - } - } finally { - db.close(); - } - const result = await readAntigravityUsage(dir, 0); - assert.deepStrictEqual(result.errors, []); - assert.deepStrictEqual( - result.files.flatMap((file) => file.records).map((record) => record.model), - ["gemini-3-pro", "claude-opus-4-6"], - ); - }); - - it("merges Antigravity aliases that bridge previously separate step records", async () => { - const db = new NodeSqlite.DatabaseSync(NodePath.join(dir, "bridge.db")); - try { - db.exec( - "CREATE TABLE gen_metadata (idx INTEGER, data BLOB); CREATE TABLE steps (idx INTEGER, metadata BLOB)", - ); - const step = db.prepare("INSERT INTO steps VALUES (?, ?)"); - step.run( - 0, - new Uint8Array(protoBytes(9, [...protoNumber(2, 100), ...protoText(11, "response")])), - ); - step.run( - 1, - new Uint8Array(protoBytes(9, [...protoNumber(3, 40), ...protoText(12, "provider")])), - ); - db.prepare("INSERT INTO gen_metadata VALUES (?, ?)").run( - 0, - new Uint8Array( - protoBytes(1, [ - ...protoText(19, "Gemini 3 Pro"), - ...protoBytes(4, [ - ...protoNumber(2, 50), - ...protoNumber(5, 20), - ...protoText(11, "response"), - ...protoText(12, "provider"), - ]), - ]), - ), - ); - } finally { - db.close(); - } - const result = await readAntigravityUsage(dir, 0); - assert.deepStrictEqual(result.errors, []); - const records = result.files.flatMap((file) => file.records); - assert.strictEqual(records.length, 1); - assert.deepStrictEqual(records[0]?.totals, { - uncachedInputTokens: 100, - cachedInputTokens: 20, - cacheCreationTokens: 0, - outputTokens: 40, - reasoningTokens: 0, - }); - }); - - it("merges Antigravity provider and message aliases across configured roots while keeping original ownership", async () => { - const roots = [NodePath.join(dir, "first"), NodePath.join(dir, "second")]; - for (const [index, root] of roots.entries()) { - await NodeFSP.mkdir(root); - const db = new NodeSqlite.DatabaseSync(NodePath.join(root, `session-${index}.db`)); - try { - db.exec("CREATE TABLE steps (idx INTEGER, metadata BLOB)"); - for (const identity of [7, 12]) { - const usage = [ - ...protoNumber(1, 246), - ...protoNumber(2, index === 0 ? 100 : 150), - ...protoText(11, `response-${index}-${identity}`), - ...protoText(identity, `shared-${identity}`), - ]; - db.prepare("INSERT INTO steps VALUES (?, ?)").run( - identity, - new Uint8Array(protoBytes(9, usage)), - ); - } - } finally { - db.close(); - } - } - const result = await readAntigravityUsage(roots, 0); - assert.deepStrictEqual(result.errors, []); - assert.strictEqual(result.files.length, 2); - assert.strictEqual(result.files[0]?.root, roots[0]); - assert.strictEqual(result.files[0]?.records.length, 2); - assert.strictEqual(result.files[1]?.records.length, 0); - assert.deepStrictEqual( - result.files[0]?.records.map((record) => record.totals.uncachedInputTokens), - [150, 150], - ); - assert.isTrue(result.files[0]?.records.every((record) => record.sessionId === "session-0")); - }); - - it("upgrades Antigravity fallback timestamps before applying the date window", async () => { - for (const fallback of ["mtime", "trajectory"]) { - const path = NodePath.join(dir, `${fallback}.db`); - const db = new NodeSqlite.DatabaseSync(path); - try { - db.exec( - "CREATE TABLE gen_metadata (idx INTEGER, data BLOB); CREATE TABLE steps (idx INTEGER, metadata BLOB)", - ); - if (fallback === "trajectory") { - db.exec("CREATE TABLE trajectory_metadata_blob (data BLOB)"); - db.prepare("INSERT INTO trajectory_metadata_blob VALUES (?)").run( - new Uint8Array(protoBytes(2, protoNumber(1, 1780000200))), - ); - } - for (const [index, seconds] of [1780000000, 1780000200].entries()) { - const usage = [...protoNumber(2, 10), ...protoText(11, `${fallback}-${index}`)]; - db.prepare("INSERT INTO steps VALUES (?, ?)").run( - index, - new Uint8Array(protoBytes(9, usage)), - ); - db.prepare("INSERT INTO gen_metadata VALUES (?, ?)").run( - index, - new Uint8Array( - protoBytes(1, [ - ...protoBytes(4, usage), - ...protoBytes(9, protoBytes(4, protoNumber(1, seconds))), - ]), - ), - ); - } - } finally { - db.close(); - } - await NodeFSP.utimes(path, 1780000000, 1780000000); - } - const result = await readAntigravityUsage(dir, 1780000100000); - assert.deepStrictEqual(result.errors, []); - const records = result.files.flatMap((file) => file.records); - assert.strictEqual(records.length, 2); - assert.deepStrictEqual( - records.map((record) => record.timestampMs), - [1780000200000, 1780000200000], - ); - }); - - const antigravityGeneration = (responseId: string) => - new Uint8Array( - protoBytes(1, [ - ...protoBytes(4, [ - ...protoNumber(2, 100), - ...protoNumber(3, 40), - ...protoText(11, responseId), - ]), - ...protoText(19, "Gemini 3 Pro"), - ...protoBytes(9, protoBytes(4, protoNumber(1, 1780000000))), - ]), - ); - - it("reuses an unchanged Antigravity database instead of decoding it again", async () => { - const path = NodePath.join(dir, "session-1.db"); - const db = new NodeSqlite.DatabaseSync(path); - try { - db.exec("CREATE TABLE gen_metadata (idx INTEGER, data BLOB)"); - db.prepare("INSERT INTO gen_metadata VALUES (?, ?)").run(0, antigravityGeneration("r-1")); - const cache = makeAntigravityUsageCache(); - assert.deepStrictEqual((await readAntigravityUsage(dir, 0, cache)).errors, []); - - // An exclusive lock makes a fresh read fail without touching the file, so - // only a cache hit can still return the earlier records. - db.exec("BEGIN EXCLUSIVE"); - assert.strictEqual((await readAntigravityUsage(dir, 0)).errors.length, 1); - const cached = await readAntigravityUsage(dir, 0, cache); - assert.deepStrictEqual(cached.errors, []); - assert.deepStrictEqual( - cached.files.flatMap((file) => file.records).map((record) => record.dedupeKey), - ["antigravity:11:r-1"], - ); - db.exec("ROLLBACK"); - } finally { - db.close(); - } - }); - - it("rereads an Antigravity database rewritten with its size and mtime restored", async () => { - const path = NodePath.join(dir, "session-1.db"); - const db = new NodeSqlite.DatabaseSync(path); - try { - db.exec("CREATE TABLE gen_metadata (idx INTEGER, data BLOB)"); - db.prepare("INSERT INTO gen_metadata VALUES (?, ?)").run(0, antigravityGeneration("r-1")); - } finally { - db.close(); - } - await NodeFSP.utimes(path, 1780000000, 1780000000); - const cache = makeAntigravityUsageCache(); - assert.deepStrictEqual((await readAntigravityUsage(dir, 0, cache)).errors, []); - - // ctime has the kernel's timestamp granularity, which can be a few - // milliseconds, so repeat the forged rewrite until it lands on a later tick - // than the cached read, as any real rewrite does. - const cached = await NodeFSP.stat(path); - do { - await NodeFSP.writeFile(path, Buffer.alloc(cached.size)); - await NodeFSP.utimes(path, 1780000000, 1780000000); - } while ((await NodeFSP.stat(path)).ctimeMs === cached.ctimeMs); - const restored = await NodeFSP.stat(path); - assert.strictEqual(restored.size, cached.size); - assert.strictEqual(restored.mtimeMs, cached.mtimeMs); - const next = await readAntigravityUsage(dir, 0, cache); - assert.deepStrictEqual(next.errors, [path]); - assert.deepStrictEqual( - next.files.flatMap((file) => file.records), - [], - ); - }); - - it("rereads an Antigravity database when only its WAL changed", async () => { - const path = NodePath.join(dir, "session-1.db"); - const db = new NodeSqlite.DatabaseSync(path); - try { - db.exec( - "PRAGMA journal_mode = WAL; PRAGMA wal_autocheckpoint = 0; CREATE TABLE gen_metadata (idx INTEGER, data BLOB)", - ); - const insert = db.prepare("INSERT INTO gen_metadata VALUES (?, ?)"); - insert.run(0, antigravityGeneration("r-1")); - const cache = makeAntigravityUsageCache(); - const first = await readAntigravityUsage(dir, 0, cache); - assert.strictEqual(first.files.flatMap((file) => file.records).length, 1); - - const before = await NodeFSP.stat(path); - insert.run(1, antigravityGeneration("r-2")); - const after = await NodeFSP.stat(path); - assert.strictEqual(after.size, before.size); - assert.strictEqual(after.mtimeMs, before.mtimeMs); - - const next = await readAntigravityUsage(dir, 0, cache); - assert.deepStrictEqual( - next.files.flatMap((file) => file.records).map((record) => record.dedupeKey), - ["antigravity:11:r-1", "antigravity:11:r-2"], - ); - } finally { - db.close(); - } - }); - - it("reads Antigravity step-only stores and reports malformed databases", async () => { - const db = new NodeSqlite.DatabaseSync(NodePath.join(dir, "steps.db")); - try { - db.exec("CREATE TABLE steps (idx INTEGER, metadata BLOB)"); - const usage = [...protoNumber(1, 246), ...protoNumber(2, 10), ...protoNumber(3, 5)]; - db.prepare("INSERT INTO steps VALUES (?, ?)").run( - 0, - new Uint8Array([...protoBytes(9, usage), ...protoBytes(8, protoNumber(1, 1780000000))]), - ); - } finally { - db.close(); - } - await NodeFSP.writeFile(NodePath.join(dir, "broken.db"), "not a sqlite database"); - const result = await readAntigravityUsage(dir, 0); - assert.strictEqual(result.errors.length, 1); - assert.strictEqual(result.files.flatMap((file) => file.records)[0]?.model, "gemini-2.5-pro"); - assert.strictEqual(result.files.flatMap((file) => file.records)[0]?.totals.outputTokens, 5); - }); - - it("ignores large values in unused Antigravity protobuf fields", async () => { - const db = new NodeSqlite.DatabaseSync(NodePath.join(dir, "large-varint.db")); - try { - db.exec("CREATE TABLE steps (idx INTEGER, metadata BLOB)"); - const unusedField = [...protoNumber(99, 0).slice(0, -1), ...Array(9).fill(0xff), 0x01]; - const usage = [...protoNumber(1, 246), ...protoNumber(2, 10), ...unusedField]; - db.prepare("INSERT INTO steps VALUES (?, ?)").run(0, new Uint8Array(protoBytes(9, usage))); - } finally { - db.close(); - } - const result = await readAntigravityUsage(dir, 0); - assert.deepStrictEqual(result.errors, []); - assert.strictEqual( - result.files.flatMap((file) => file.records)[0]?.totals.uncachedInputTokens, - 10, - ); - }); -}); diff --git a/packages/provider-core/src/server/usage.ts b/packages/provider-core/src/server/usage.ts index 62ecb5796b1b0dffed7fa08deea3996a660f9d81..0f1235d7037d8d6a719a4b998b70d450d8017fd4 100644 --- a/packages/provider-core/src/server/usage.ts +++ b/packages/provider-core/src/server/usage.ts @@ -13,6 +13,7 @@ import type { ProviderInstanceId, ServerSettings, UsageProviderKind, + UsageReadError, UsageSource, UsageTokenTotals, } from "@t3tools/contracts"; @@ -177,7 +178,8 @@ export type ProviderUsageReader = | { readonly kind: "scan"; readonly provider: UsageProviderKind; + /** Fails only when the whole usage read cannot go on; an unreadable source is a status. */ readonly scan: ( input: ProviderUsageScanInput, - ) => Effect.Effect, never, R>; + ) => Effect.Effect, UsageReadError, R>; };