| 1 | // @effect-diagnostics-next-line nodeBuiltinImport:off - Effect's File.Info has no ctime or sub-millisecond mtime, which the cache tests compare. |
| 2 | import * as NodeFSP from "node:fs/promises"; |
| 3 | import * as NodeSqlite from "node:sqlite"; |
| 4 | |
| 5 | import * as NodeServices from "@effect/platform-node/NodeServices"; |
| 6 | import { assert, describe, it, vi } from "@effect/vitest"; |
| 7 | import * as NodeSqliteClient from "@t3tools/shared/nodeSqliteClient"; |
| 8 | import * as Context from "effect/Context"; |
| 9 | import * as Effect from "effect/Effect"; |
| 10 | import * as FileSystem from "effect/FileSystem"; |
| 11 | import * as Layer from "effect/Layer"; |
| 12 | import * as Path from "effect/Path"; |
| 13 | import * as SqlClient from "effect/sql/SqlClient"; |
| 14 | import * as Tracer from "effect/Tracer"; |
| 15 | |
| 16 | import { makeAntigravityUsageCache, readAntigravityUsage } from "./antigravityUsageReader.ts"; |
| 17 | |
| 18 | /** A writable connection that stays open until the enclosing scope closes. */ |
| 19 | const openDatabase = Effect.fn("openDatabase")(function* (filename: string) { |
| 20 | const context = yield* Layer.build(NodeSqliteClient.layer({ filename })); |
| 21 | const sql = Context.get(context, SqlClient.SqlClient); |
| 22 | return { |
| 23 | exec: (statement: string) => sql.unsafe(statement), |
| 24 | run: (statement: string, params: ReadonlyArray<unknown>) => sql.unsafe(statement, params), |
| 25 | }; |
| 26 | }); |
| 27 | |
| 28 | /** Seeds `filename` and closes it, as a provider that has finished writing. */ |
| 29 | const seedDatabase = <E>( |
| 30 | filename: string, |
| 31 | seed: (db: Effect.Success<ReturnType<typeof openDatabase>>) => Effect.Effect<void, E>, |
| 32 | ) => |
| 33 | Effect.scoped( |
| 34 | Effect.gen(function* () { |
| 35 | yield* seed(yield* openDatabase(filename)); |
| 36 | }), |
| 37 | ); |
| 38 | |
| 39 | const tempDir = Effect.gen(function* () { |
| 40 | const fileSystem = yield* FileSystem.FileSystem; |
| 41 | return yield* fileSystem.makeTempDirectoryScoped({ prefix: "usage-reader-test-" }); |
| 42 | }); |
| 43 | |
| 44 | function protoNumber(field: number, value: number): number[] { |
| 45 | const varint = (number: number) => { |
| 46 | const bytes: number[] = []; |
| 47 | do { |
| 48 | const byte = number % 128; |
| 49 | number = Math.floor(number / 128); |
| 50 | bytes.push(byte + (number > 0 ? 128 : 0)); |
| 51 | } while (number > 0); |
| 52 | return bytes; |
| 53 | }; |
| 54 | return [...varint(field * 8), ...varint(value)]; |
| 55 | } |
| 56 | |
| 57 | function protoBytes(field: number, bytes: readonly number[]): number[] { |
| 58 | const encoded = protoNumber(field, bytes.length); |
| 59 | encoded[0] = encoded[0]! + 2; |
| 60 | return [...encoded, ...bytes]; |
| 61 | } |
| 62 | |
| 63 | function protoText(field: number, value: string): number[] { |
| 64 | return protoBytes(field, [...Buffer.from(value)]); |
| 65 | } |
| 66 | |
| 67 | const TABLES = |
| 68 | "CREATE TABLE gen_metadata (idx INTEGER, data BLOB); CREATE TABLE steps (idx INTEGER, metadata BLOB)"; |
| 69 | |
| 70 | const createTables = (db: Effect.Success<ReturnType<typeof openDatabase>>) => |
| 71 | Effect.forEach(TABLES.split("; "), db.exec, { discard: true }); |
| 72 | |
| 73 | const records = (result: Effect.Success<ReturnType<typeof readAntigravityUsage>>) => |
| 74 | result.files.flatMap((file) => file.records); |
| 75 | |
| 76 | describe("readAntigravityUsage", () => { |
| 77 | it.effect( |
| 78 | "deduplicates Antigravity generation and step usage while preserving retry model and token buckets", |
| 79 | () => |
| 80 | Effect.gen(function* () { |
| 81 | const path = yield* Path.Path; |
| 82 | const dir = yield* tempDir; |
| 83 | const stamp = protoNumber(1, 1780000000); |
| 84 | const usage = [ |
| 85 | ...protoNumber(2, 100), |
| 86 | ...protoNumber(3, 40), |
| 87 | ...protoNumber(4, 5), |
| 88 | ...protoNumber(5, 20), |
| 89 | ...protoNumber(9, 10), |
| 90 | ...protoText(11, "response-1"), |
| 91 | ]; |
| 92 | const retry = [ |
| 93 | ...protoNumber(1, 1026), |
| 94 | ...protoNumber(2, 12), |
| 95 | ...protoNumber(3, 3), |
| 96 | ...protoText(11, "retry-1"), |
| 97 | ]; |
| 98 | const generation = protoBytes(1, [ |
| 99 | ...protoBytes(4, usage), |
| 100 | ...protoText(19, "Gemini 3 Pro"), |
| 101 | ...protoBytes(9, protoBytes(4, stamp)), |
| 102 | ]); |
| 103 | const step = [ |
| 104 | ...protoBytes(9, usage), |
| 105 | ...protoBytes(8, stamp), |
| 106 | ...protoBytes(28, protoBytes(2, retry)), |
| 107 | ]; |
| 108 | yield* seedDatabase(path.join(dir, "session-1.db"), (db) => |
| 109 | Effect.gen(function* () { |
| 110 | yield* createTables(db); |
| 111 | yield* db.run("INSERT INTO gen_metadata VALUES (?, ?)", [ |
| 112 | 0, |
| 113 | new Uint8Array(generation), |
| 114 | ]); |
| 115 | yield* db.run("INSERT INTO steps VALUES (?, ?)", [0, new Uint8Array(step)]); |
| 116 | }), |
| 117 | ); |
| 118 | const result = yield* readAntigravityUsage(dir, 0); |
| 119 | assert.deepStrictEqual(result.errors, []); |
| 120 | const read = records(result); |
| 121 | assert.strictEqual(read.length, 2); |
| 122 | const main = read.find((record) => record.model === "gemini-3-pro"); |
| 123 | assert.isDefined(main); |
| 124 | assert.strictEqual(main?.timestampMs, 1780000000000); |
| 125 | assert.strictEqual(main?.sessionId, "session-1"); |
| 126 | assert.deepStrictEqual(main?.totals, { |
| 127 | uncachedInputTokens: 100, |
| 128 | cachedInputTokens: 20, |
| 129 | cacheCreationTokens: 5, |
| 130 | outputTokens: 40, |
| 131 | reasoningTokens: 10, |
| 132 | }); |
| 133 | assert.strictEqual( |
| 134 | read.find((record) => record.model === "claude-opus-4-6")?.totals.uncachedInputTokens, |
| 135 | 12, |
| 136 | ); |
| 137 | assert.deepStrictEqual(records(yield* readAntigravityUsage(dir, 1780000000001)), []); |
| 138 | }).pipe(Effect.provide(NodeServices.layer)), |
| 139 | ); |
| 140 | |
| 141 | it.effect("uses the matching Antigravity generation model for each model-less step", () => |
| 142 | Effect.gen(function* () { |
| 143 | const path = yield* Path.Path; |
| 144 | const dir = yield* tempDir; |
| 145 | yield* seedDatabase(path.join(dir, "model-switch.db"), (db) => |
| 146 | Effect.gen(function* () { |
| 147 | yield* createTables(db); |
| 148 | for (const [idx, name] of ["Gemini 3 Pro", "Claude Opus 4.6"].entries()) { |
| 149 | yield* db.run("INSERT INTO gen_metadata VALUES (?, ?)", [ |
| 150 | idx, |
| 151 | new Uint8Array(protoBytes(1, protoText(19, name))), |
| 152 | ]); |
| 153 | yield* db.run("INSERT INTO steps VALUES (?, ?)", [ |
| 154 | idx, |
| 155 | new Uint8Array(protoBytes(9, protoNumber(2, 10 + idx))), |
| 156 | ]); |
| 157 | } |
| 158 | }), |
| 159 | ); |
| 160 | const result = yield* readAntigravityUsage(dir, 0); |
| 161 | assert.deepStrictEqual(result.errors, []); |
| 162 | assert.deepStrictEqual( |
| 163 | records(result).map((record) => record.model), |
| 164 | ["gemini-3-pro", "claude-opus-4-6"], |
| 165 | ); |
| 166 | }).pipe(Effect.provide(NodeServices.layer)), |
| 167 | ); |
| 168 | |
| 169 | it.effect("merges Antigravity aliases that bridge previously separate step records", () => |
| 170 | Effect.gen(function* () { |
| 171 | const path = yield* Path.Path; |
| 172 | const dir = yield* tempDir; |
| 173 | yield* seedDatabase(path.join(dir, "bridge.db"), (db) => |
| 174 | Effect.gen(function* () { |
| 175 | yield* createTables(db); |
| 176 | yield* db.run("INSERT INTO steps VALUES (?, ?)", [ |
| 177 | 0, |
| 178 | new Uint8Array(protoBytes(9, [...protoNumber(2, 100), ...protoText(11, "response")])), |
| 179 | ]); |
| 180 | yield* db.run("INSERT INTO steps VALUES (?, ?)", [ |
| 181 | 1, |
| 182 | new Uint8Array(protoBytes(9, [...protoNumber(3, 40), ...protoText(12, "provider")])), |
| 183 | ]); |
| 184 | yield* db.run("INSERT INTO gen_metadata VALUES (?, ?)", [ |
| 185 | 0, |
| 186 | new Uint8Array( |
| 187 | protoBytes(1, [ |
| 188 | ...protoText(19, "Gemini 3 Pro"), |
| 189 | ...protoBytes(4, [ |
| 190 | ...protoNumber(2, 50), |
| 191 | ...protoNumber(5, 20), |
| 192 | ...protoText(11, "response"), |
| 193 | ...protoText(12, "provider"), |
| 194 | ]), |
| 195 | ]), |
| 196 | ), |
| 197 | ]); |
| 198 | }), |
| 199 | ); |
| 200 | const result = yield* readAntigravityUsage(dir, 0); |
| 201 | assert.deepStrictEqual(result.errors, []); |
| 202 | const read = records(result); |
| 203 | assert.strictEqual(read.length, 1); |
| 204 | assert.deepStrictEqual(read[0]?.totals, { |
| 205 | uncachedInputTokens: 100, |
| 206 | cachedInputTokens: 20, |
| 207 | cacheCreationTokens: 0, |
| 208 | outputTokens: 40, |
| 209 | reasoningTokens: 0, |
| 210 | }); |
| 211 | }).pipe(Effect.provide(NodeServices.layer)), |
| 212 | ); |
| 213 | |
| 214 | it.effect( |
| 215 | "merges Antigravity provider and message aliases across configured roots while keeping original ownership", |
| 216 | () => |
| 217 | Effect.gen(function* () { |
| 218 | const fileSystem = yield* FileSystem.FileSystem; |
| 219 | const path = yield* Path.Path; |
| 220 | const dir = yield* tempDir; |
| 221 | const roots = [path.join(dir, "first"), path.join(dir, "second")]; |
| 222 | for (const [index, root] of roots.entries()) { |
| 223 | yield* fileSystem.makeDirectory(root); |
| 224 | yield* seedDatabase(path.join(root, `session-${index}.db`), (db) => |
| 225 | Effect.gen(function* () { |
| 226 | yield* db.exec("CREATE TABLE steps (idx INTEGER, metadata BLOB)"); |
| 227 | for (const identity of [7, 12]) { |
| 228 | const usage = [ |
| 229 | ...protoNumber(1, 246), |
| 230 | ...protoNumber(2, index === 0 ? 100 : 150), |
| 231 | ...protoText(11, `response-${index}-${identity}`), |
| 232 | ...protoText(identity, `shared-${identity}`), |
| 233 | ]; |
| 234 | yield* db.run("INSERT INTO steps VALUES (?, ?)", [ |
| 235 | identity, |
| 236 | new Uint8Array(protoBytes(9, usage)), |
| 237 | ]); |
| 238 | } |
| 239 | }), |
| 240 | ); |
| 241 | } |
| 242 | const result = yield* readAntigravityUsage(roots, 0); |
| 243 | assert.deepStrictEqual(result.errors, []); |
| 244 | assert.strictEqual(result.files.length, 2); |
| 245 | assert.strictEqual(result.files[0]?.root, roots[0]); |
| 246 | assert.strictEqual(result.files[0]?.records.length, 2); |
| 247 | assert.strictEqual(result.files[1]?.records.length, 0); |
| 248 | assert.deepStrictEqual( |
| 249 | result.files[0]?.records.map((record) => record.totals.uncachedInputTokens), |
| 250 | [150, 150], |
| 251 | ); |
| 252 | assert.isTrue(result.files[0]?.records.every((record) => record.sessionId === "session-0")); |
| 253 | }).pipe(Effect.provide(NodeServices.layer)), |
| 254 | ); |
| 255 | |
| 256 | it.effect("upgrades Antigravity fallback timestamps before applying the date window", () => |
| 257 | Effect.gen(function* () { |
| 258 | const fileSystem = yield* FileSystem.FileSystem; |
| 259 | const path = yield* Path.Path; |
| 260 | const dir = yield* tempDir; |
| 261 | for (const fallback of ["mtime", "trajectory"]) { |
| 262 | const file = path.join(dir, `${fallback}.db`); |
| 263 | yield* seedDatabase(file, (db) => |
| 264 | Effect.gen(function* () { |
| 265 | yield* createTables(db); |
| 266 | if (fallback === "trajectory") { |
| 267 | yield* db.exec("CREATE TABLE trajectory_metadata_blob (data BLOB)"); |
| 268 | yield* db.run("INSERT INTO trajectory_metadata_blob VALUES (?)", [ |
| 269 | new Uint8Array(protoBytes(2, protoNumber(1, 1780000200))), |
| 270 | ]); |
| 271 | } |
| 272 | for (const [index, seconds] of [1780000000, 1780000200].entries()) { |
| 273 | const usage = [...protoNumber(2, 10), ...protoText(11, `${fallback}-${index}`)]; |
| 274 | yield* db.run("INSERT INTO steps VALUES (?, ?)", [ |
| 275 | index, |
| 276 | new Uint8Array(protoBytes(9, usage)), |
| 277 | ]); |
| 278 | yield* db.run("INSERT INTO gen_metadata VALUES (?, ?)", [ |
| 279 | index, |
| 280 | new Uint8Array( |
| 281 | protoBytes(1, [ |
| 282 | ...protoBytes(4, usage), |
| 283 | ...protoBytes(9, protoBytes(4, protoNumber(1, seconds))), |
| 284 | ]), |
| 285 | ), |
| 286 | ]); |
| 287 | } |
| 288 | }), |
| 289 | ); |
| 290 | yield* fileSystem.utimes(file, 1780000000, 1780000000); |
| 291 | } |
| 292 | const result = yield* readAntigravityUsage(dir, 1780000100000); |
| 293 | assert.deepStrictEqual(result.errors, []); |
| 294 | const read = records(result); |
| 295 | assert.strictEqual(read.length, 2); |
| 296 | assert.deepStrictEqual( |
| 297 | read.map((record) => record.timestampMs), |
| 298 | [1780000200000, 1780000200000], |
| 299 | ); |
| 300 | }).pipe(Effect.provide(NodeServices.layer)), |
| 301 | ); |
| 302 | |
| 303 | const antigravityGeneration = (responseId: string) => |
| 304 | new Uint8Array( |
| 305 | protoBytes(1, [ |
| 306 | ...protoBytes(4, [ |
| 307 | ...protoNumber(2, 100), |
| 308 | ...protoNumber(3, 40), |
| 309 | ...protoText(11, responseId), |
| 310 | ]), |
| 311 | ...protoText(19, "Gemini 3 Pro"), |
| 312 | ...protoBytes(9, protoBytes(4, protoNumber(1, 1780000000))), |
| 313 | ]), |
| 314 | ); |
| 315 | |
| 316 | it.effect("reuses an unchanged Antigravity database instead of decoding it again", () => |
| 317 | Effect.gen(function* () { |
| 318 | const path = yield* Path.Path; |
| 319 | const dir = yield* tempDir; |
| 320 | const db = yield* openDatabase(path.join(dir, "session-1.db")); |
| 321 | yield* db.exec("CREATE TABLE gen_metadata (idx INTEGER, data BLOB)"); |
| 322 | yield* db.run("INSERT INTO gen_metadata VALUES (?, ?)", [0, antigravityGeneration("r-1")]); |
| 323 | const cache = makeAntigravityUsageCache(); |
| 324 | assert.deepStrictEqual((yield* readAntigravityUsage(dir, 0, cache)).errors, []); |
| 325 | |
| 326 | // An exclusive lock makes a fresh read fail without touching the file, so |
| 327 | // only a cache hit can still return the earlier records. |
| 328 | yield* db.exec("BEGIN EXCLUSIVE"); |
| 329 | assert.strictEqual((yield* readAntigravityUsage(dir, 0)).errors.length, 1); |
| 330 | const cached = yield* readAntigravityUsage(dir, 0, cache); |
| 331 | assert.deepStrictEqual(cached.errors, []); |
| 332 | assert.deepStrictEqual( |
| 333 | records(cached).map((record) => record.dedupeKey), |
| 334 | ["antigravity:11:r-1"], |
| 335 | ); |
| 336 | yield* db.exec("ROLLBACK"); |
| 337 | }).pipe(Effect.provide(NodeServices.layer)), |
| 338 | ); |
| 339 | |
| 340 | it.effect("rereads an Antigravity database rewritten with its size and mtime restored", () => |
| 341 | Effect.gen(function* () { |
| 342 | const fileSystem = yield* FileSystem.FileSystem; |
| 343 | const pathService = yield* Path.Path; |
| 344 | const dir = yield* tempDir; |
| 345 | const path = pathService.join(dir, "session-1.db"); |
| 346 | yield* seedDatabase(path, (db) => |
| 347 | Effect.gen(function* () { |
| 348 | yield* db.exec("CREATE TABLE gen_metadata (idx INTEGER, data BLOB)"); |
| 349 | yield* db.run("INSERT INTO gen_metadata VALUES (?, ?)", [ |
| 350 | 0, |
| 351 | antigravityGeneration("r-1"), |
| 352 | ]); |
| 353 | }), |
| 354 | ); |
| 355 | yield* fileSystem.utimes(path, 1780000000, 1780000000); |
| 356 | const cache = makeAntigravityUsageCache(); |
| 357 | assert.deepStrictEqual((yield* readAntigravityUsage(dir, 0, cache)).errors, []); |
| 358 | |
| 359 | // ctime has the kernel's timestamp granularity, which can be a few |
| 360 | // milliseconds, so repeat the forged rewrite until it lands on a later tick |
| 361 | // than the cached read, as any real rewrite does. |
| 362 | const stat = () => Effect.promise(() => NodeFSP.stat(path)); |
| 363 | const cached = yield* stat(); |
| 364 | do { |
| 365 | yield* fileSystem.writeFile(path, new Uint8Array(cached.size)); |
| 366 | yield* fileSystem.utimes(path, 1780000000, 1780000000); |
| 367 | } while ((yield* stat()).ctimeMs === cached.ctimeMs); |
| 368 | const restored = yield* stat(); |
| 369 | assert.strictEqual(restored.size, cached.size); |
| 370 | assert.strictEqual(restored.mtimeMs, cached.mtimeMs); |
| 371 | const next = yield* readAntigravityUsage(dir, 0, cache); |
| 372 | assert.deepStrictEqual(next.errors, [path]); |
| 373 | assert.deepStrictEqual(records(next), []); |
| 374 | }).pipe(Effect.provide(NodeServices.layer)), |
| 375 | ); |
| 376 | |
| 377 | it.effect("rereads an Antigravity database when only its WAL changed", () => |
| 378 | Effect.gen(function* () { |
| 379 | const pathService = yield* Path.Path; |
| 380 | const dir = yield* tempDir; |
| 381 | const path = pathService.join(dir, "session-1.db"); |
| 382 | const db = yield* openDatabase(path); |
| 383 | yield* db.exec("PRAGMA journal_mode = WAL"); |
| 384 | yield* db.exec("PRAGMA wal_autocheckpoint = 0"); |
| 385 | yield* db.exec("CREATE TABLE gen_metadata (idx INTEGER, data BLOB)"); |
| 386 | const insert = (idx: number, responseId: string) => |
| 387 | db.run("INSERT INTO gen_metadata VALUES (?, ?)", [idx, antigravityGeneration(responseId)]); |
| 388 | yield* insert(0, "r-1"); |
| 389 | const cache = makeAntigravityUsageCache(); |
| 390 | const first = yield* readAntigravityUsage(dir, 0, cache); |
| 391 | assert.strictEqual(records(first).length, 1); |
| 392 | |
| 393 | const before = yield* Effect.promise(() => NodeFSP.stat(path)); |
| 394 | yield* insert(1, "r-2"); |
| 395 | const after = yield* Effect.promise(() => NodeFSP.stat(path)); |
| 396 | assert.strictEqual(after.size, before.size); |
| 397 | assert.strictEqual(after.mtimeMs, before.mtimeMs); |
| 398 | |
| 399 | const next = yield* readAntigravityUsage(dir, 0, cache); |
| 400 | assert.deepStrictEqual( |
| 401 | records(next).map((record) => record.dedupeKey), |
| 402 | ["antigravity:11:r-1", "antigravity:11:r-2"], |
| 403 | ); |
| 404 | }).pipe(Effect.provide(NodeServices.layer)), |
| 405 | ); |
| 406 | |
| 407 | it.effect("reads Antigravity step-only stores and reports malformed databases", () => |
| 408 | Effect.gen(function* () { |
| 409 | const fileSystem = yield* FileSystem.FileSystem; |
| 410 | const path = yield* Path.Path; |
| 411 | const dir = yield* tempDir; |
| 412 | yield* seedDatabase(path.join(dir, "steps.db"), (db) => |
| 413 | Effect.gen(function* () { |
| 414 | yield* db.exec("CREATE TABLE steps (idx INTEGER, metadata BLOB)"); |
| 415 | const usage = [...protoNumber(1, 246), ...protoNumber(2, 10), ...protoNumber(3, 5)]; |
| 416 | yield* db.run("INSERT INTO steps VALUES (?, ?)", [ |
| 417 | 0, |
| 418 | new Uint8Array([...protoBytes(9, usage), ...protoBytes(8, protoNumber(1, 1780000000))]), |
| 419 | ]); |
| 420 | }), |
| 421 | ); |
| 422 | yield* fileSystem.writeFileString(path.join(dir, "broken.db"), "not a sqlite database"); |
| 423 | const result = yield* readAntigravityUsage(dir, 0); |
| 424 | assert.strictEqual(result.errors.length, 1); |
| 425 | assert.strictEqual(records(result)[0]?.model, "gemini-2.5-pro"); |
| 426 | assert.strictEqual(records(result)[0]?.totals.outputTokens, 5); |
| 427 | }).pipe(Effect.provide(NodeServices.layer)), |
| 428 | ); |
| 429 | |
| 430 | it.effect("ignores large values in unused Antigravity protobuf fields", () => |
| 431 | Effect.gen(function* () { |
| 432 | const path = yield* Path.Path; |
| 433 | const dir = yield* tempDir; |
| 434 | yield* seedDatabase(path.join(dir, "large-varint.db"), (db) => |
| 435 | Effect.gen(function* () { |
| 436 | yield* db.exec("CREATE TABLE steps (idx INTEGER, metadata BLOB)"); |
| 437 | const unusedField = [...protoNumber(99, 0).slice(0, -1), ...Array(9).fill(0xff), 0x01]; |
| 438 | const usage = [...protoNumber(1, 246), ...protoNumber(2, 10), ...unusedField]; |
| 439 | yield* db.run("INSERT INTO steps VALUES (?, ?)", [ |
| 440 | 0, |
| 441 | new Uint8Array(protoBytes(9, usage)), |
| 442 | ]); |
| 443 | }), |
| 444 | ); |
| 445 | const result = yield* readAntigravityUsage(dir, 0); |
| 446 | assert.deepStrictEqual(result.errors, []); |
| 447 | assert.strictEqual(records(result)[0]?.totals.uncachedInputTokens, 10); |
| 448 | }).pipe(Effect.provide(NodeServices.layer)), |
| 449 | ); |
| 450 | |
| 451 | /** Counts the rows each StatementSync.all() call returns, by its SQL. */ |
| 452 | const recordReads = Effect.gen(function* () { |
| 453 | const reads: Array<{ sql: string; rows: number }> = []; |
| 454 | const all = NodeSqlite.StatementSync.prototype.all; |
| 455 | yield* Effect.acquireRelease( |
| 456 | Effect.sync(() => |
| 457 | vi.spyOn(NodeSqlite.StatementSync.prototype, "all").mockImplementation(function ( |
| 458 | this: NodeSqlite.StatementSync, |
| 459 | ...params |
| 460 | ) { |
| 461 | const result = all.apply(this, params); |
| 462 | reads.push({ sql: this.sourceSQL, rows: result.length }); |
| 463 | return result; |
| 464 | }), |
| 465 | ), |
| 466 | (spy) => Effect.sync(() => spy.mockRestore()), |
| 467 | ); |
| 468 | return reads; |
| 469 | }); |
| 470 | |
| 471 | const stepUsage = (responseId: string) => |
| 472 | new Uint8Array(protoBytes(9, [...protoNumber(2, 10), ...protoText(11, responseId)])); |
| 473 | |
| 474 | it.effect("reads one Antigravity generation or one page of steps per statement", () => |
| 475 | Effect.gen(function* () { |
| 476 | const path = yield* Path.Path; |
| 477 | const dir = yield* tempDir; |
| 478 | yield* seedDatabase(path.join(dir, "session-1.db"), (db) => |
| 479 | Effect.gen(function* () { |
| 480 | // Antigravity's keys and indexes: idx is the INTEGER PRIMARY KEY, and steps has indexes. |
| 481 | yield* db.exec( |
| 482 | "CREATE TABLE `gen_metadata` (`idx` integer,`data` blob,`size` integer NOT NULL DEFAULT 0,PRIMARY KEY (`idx`))", |
| 483 | ); |
| 484 | yield* db.exec( |
| 485 | "CREATE TABLE `steps` (`idx` integer,`step_type` integer NOT NULL DEFAULT 0,`status` integer NOT NULL DEFAULT 0,`metadata` blob,PRIMARY KEY (`idx`))", |
| 486 | ); |
| 487 | yield* db.exec("CREATE INDEX `idx_steps_status` ON `steps`(`status`)"); |
| 488 | yield* db.exec("CREATE INDEX `idx_steps_step_type` ON `steps`(`step_type`)"); |
| 489 | yield* db.exec("BEGIN"); |
| 490 | for (let idx = 0; idx < 3; idx++) { |
| 491 | yield* db.run("INSERT INTO gen_metadata (idx, data) VALUES (?, ?)", [ |
| 492 | idx, |
| 493 | antigravityGeneration(`g-${idx}`), |
| 494 | ]); |
| 495 | } |
| 496 | // A run of steps without metadata sits between steps with usage. |
| 497 | for (let idx = 0; idx < 250; idx++) { |
| 498 | yield* db.run("INSERT INTO steps (idx, metadata) VALUES (?, ?)", [ |
| 499 | idx, |
| 500 | idx >= 75 && idx < 175 ? null : stepUsage(`s-${idx}`), |
| 501 | ]); |
| 502 | } |
| 503 | yield* db.exec("COMMIT"); |
| 504 | }), |
| 505 | ); |
| 506 | const reads = yield* recordReads; |
| 507 | const spans: Tracer.NativeSpan[] = []; |
| 508 | const tracer = Tracer.make({ |
| 509 | span: (options) => { |
| 510 | const span = new Tracer.NativeSpan(options); |
| 511 | spans.push(span); |
| 512 | return span; |
| 513 | }, |
| 514 | }); |
| 515 | const result = yield* readAntigravityUsage(dir, 0).pipe( |
| 516 | Effect.provideService(Tracer.Tracer, tracer), |
| 517 | ); |
| 518 | assert.deepStrictEqual(result.errors, []); |
| 519 | assert.deepStrictEqual( |
| 520 | records(result).map((record) => record.dedupeKey), |
| 521 | [ |
| 522 | ...Array.from({ length: 75 }, (_, idx) => `antigravity:11:s-${idx}`), |
| 523 | ...Array.from({ length: 75 }, (_, idx) => `antigravity:11:s-${idx + 175}`), |
| 524 | ...Array.from({ length: 3 }, (_, idx) => `antigravity:11:g-${idx}`), |
| 525 | ], |
| 526 | ); |
| 527 | // Page statements are untraced, so the trace does not grow with the tables. |
| 528 | assert.isTrue(spans.some((span) => span.name === "sqlite")); |
| 529 | assert.deepStrictEqual( |
| 530 | spans.flatMap((span) => { |
| 531 | const query = span.attributes.get("db.query.text"); |
| 532 | return typeof query === "string" && query.includes("rowid > ?") ? [query] : []; |
| 533 | }), |
| 534 | [], |
| 535 | ); |
| 536 | // Every row, with metadata or not, comes back through a bounded statement. |
| 537 | for (const [table, total, most] of [ |
| 538 | ["steps", 250, 64], |
| 539 | ["gen_metadata", 3, 1], |
| 540 | ] as const) { |
| 541 | const rows = reads |
| 542 | .filter((read) => read.sql.includes(`FROM ${table} `)) |
| 543 | .map((read) => read.rows); |
| 544 | assert.strictEqual( |
| 545 | rows.reduce((sum, count) => sum + count, 0), |
| 546 | total, |
| 547 | table, |
| 548 | ); |
| 549 | assert.isAtMost(Math.max(...rows), most, table); |
| 550 | } |
| 551 | }).pipe(Effect.provide(NodeServices.layer)), |
| 552 | ); |
| 553 | |
| 554 | it.effect("returns Antigravity records in idx order however the rows are stored", () => |
| 555 | Effect.gen(function* () { |
| 556 | const path = yield* Path.Path; |
| 557 | const dir = yield* tempDir; |
| 558 | yield* seedDatabase(path.join(dir, "session-1.db"), (db) => |
| 559 | Effect.gen(function* () { |
| 560 | yield* createTables(db); |
| 561 | yield* db.exec("BEGIN"); |
| 562 | // Stored in reverse idx order, so rowid order is the opposite of idx order. |
| 563 | for (let idx = 99; idx >= 0; idx--) { |
| 564 | yield* db.run("INSERT INTO steps VALUES (?, ?)", [idx, stepUsage(`s-${idx}`)]); |
| 565 | } |
| 566 | for (const idx of [1, 0]) { |
| 567 | yield* db.run("INSERT INTO gen_metadata VALUES (?, ?)", [ |
| 568 | idx, |
| 569 | antigravityGeneration(`g-${idx}`), |
| 570 | ]); |
| 571 | } |
| 572 | yield* db.exec("COMMIT"); |
| 573 | }), |
| 574 | ); |
| 575 | const result = yield* readAntigravityUsage(dir, 0); |
| 576 | assert.deepStrictEqual(result.errors, []); |
| 577 | assert.deepStrictEqual( |
| 578 | records(result).map((record) => record.dedupeKey), |
| 579 | [ |
| 580 | ...Array.from({ length: 100 }, (_, idx) => `antigravity:11:s-${idx}`), |
| 581 | "antigravity:11:g-0", |
| 582 | "antigravity:11:g-1", |
| 583 | ], |
| 584 | ); |
| 585 | }).pipe(Effect.provide(NodeServices.layer)), |
| 586 | ); |
| 587 | }); |