From 52ffe0cae6e5e03c95037152ef4548f40daf5c2e Mon Sep 17 00:00:00 2001 From: Theo Browne Date: Wed, 14 Oct 2026 15:53:24 -0700 Subject: [PATCH] fix(usage): cold usage scans no longer stall the server (#19446) Co-authored-by: Victor Vogel <263261067+only21mil@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 --- .../Drivers/antigravityUsageReader.test.ts | 141 +++++++++++++++++- .../Drivers/antigravityUsageReader.ts | 60 ++++++-- .../src/server/usage.test.ts | 56 +++++++ .../provider-opencode/src/server/usage.ts | 61 ++++---- 4 files changed, 274 insertions(+), 44 deletions(-) diff --git a/apps/server/src/provider/Drivers/antigravityUsageReader.test.ts b/apps/server/src/provider/Drivers/antigravityUsageReader.test.ts index 20c815fc9b3250be0a51b36d8aac4dc6dc515741..d05038ca86cc8e1fc3b1d353c9acf787910f461c 100644 --- a/apps/server/src/provider/Drivers/antigravityUsageReader.test.ts +++ b/apps/server/src/provider/Drivers/antigravityUsageReader.test.ts @@ -1,8 +1,9 @@ // @effect-diagnostics-next-line nodeBuiltinImport:off - Effect's File.Info has no ctime or sub-millisecond mtime, which the cache tests compare. import * as NodeFSP from "node:fs/promises"; +import * as NodeSqlite from "node:sqlite"; import * as NodeServices from "@effect/platform-node/NodeServices"; -import { assert, describe, it } from "@effect/vitest"; +import { assert, describe, it, vi } from "@effect/vitest"; import * as NodeSqliteClient from "@t3tools/shared/nodeSqliteClient"; import * as Context from "effect/Context"; import * as Effect from "effect/Effect"; @@ -10,6 +11,7 @@ import * as FileSystem from "effect/FileSystem"; import * as Layer from "effect/Layer"; import * as Path from "effect/Path"; import * as SqlClient from "effect/sql/SqlClient"; +import * as Tracer from "effect/Tracer"; import { makeAntigravityUsageCache, readAntigravityUsage } from "./antigravityUsageReader.ts"; @@ -445,4 +447,141 @@ describe("readAntigravityUsage", () => { assert.strictEqual(records(result)[0]?.totals.uncachedInputTokens, 10); }).pipe(Effect.provide(NodeServices.layer)), ); + + /** Counts the rows each StatementSync.all() call returns, by its SQL. */ + const recordReads = Effect.gen(function* () { + const reads: Array<{ sql: string; rows: number }> = []; + const all = NodeSqlite.StatementSync.prototype.all; + yield* Effect.acquireRelease( + Effect.sync(() => + vi.spyOn(NodeSqlite.StatementSync.prototype, "all").mockImplementation(function ( + this: NodeSqlite.StatementSync, + ...params + ) { + const result = all.apply(this, params); + reads.push({ sql: this.sourceSQL, rows: result.length }); + return result; + }), + ), + (spy) => Effect.sync(() => spy.mockRestore()), + ); + return reads; + }); + + const stepUsage = (responseId: string) => + new Uint8Array(protoBytes(9, [...protoNumber(2, 10), ...protoText(11, responseId)])); + + it.effect("reads one Antigravity generation or one page of steps per statement", () => + Effect.gen(function* () { + const path = yield* Path.Path; + const dir = yield* tempDir; + yield* seedDatabase(path.join(dir, "session-1.db"), (db) => + Effect.gen(function* () { + // Antigravity's keys and indexes: idx is the INTEGER PRIMARY KEY, and steps has indexes. + yield* db.exec( + "CREATE TABLE `gen_metadata` (`idx` integer,`data` blob,`size` integer NOT NULL DEFAULT 0,PRIMARY KEY (`idx`))", + ); + yield* db.exec( + "CREATE TABLE `steps` (`idx` integer,`step_type` integer NOT NULL DEFAULT 0,`status` integer NOT NULL DEFAULT 0,`metadata` blob,PRIMARY KEY (`idx`))", + ); + yield* db.exec("CREATE INDEX `idx_steps_status` ON `steps`(`status`)"); + yield* db.exec("CREATE INDEX `idx_steps_step_type` ON `steps`(`step_type`)"); + yield* db.exec("BEGIN"); + for (let idx = 0; idx < 3; idx++) { + yield* db.run("INSERT INTO gen_metadata (idx, data) VALUES (?, ?)", [ + idx, + antigravityGeneration(`g-${idx}`), + ]); + } + // A run of steps without metadata sits between steps with usage. + for (let idx = 0; idx < 250; idx++) { + yield* db.run("INSERT INTO steps (idx, metadata) VALUES (?, ?)", [ + idx, + idx >= 75 && idx < 175 ? null : stepUsage(`s-${idx}`), + ]); + } + yield* db.exec("COMMIT"); + }), + ); + const reads = yield* recordReads; + const spans: Tracer.NativeSpan[] = []; + const tracer = Tracer.make({ + span: (options) => { + const span = new Tracer.NativeSpan(options); + spans.push(span); + return span; + }, + }); + const result = yield* readAntigravityUsage(dir, 0).pipe( + Effect.provideService(Tracer.Tracer, tracer), + ); + assert.deepStrictEqual(result.errors, []); + assert.deepStrictEqual( + records(result).map((record) => record.dedupeKey), + [ + ...Array.from({ length: 75 }, (_, idx) => `antigravity:11:s-${idx}`), + ...Array.from({ length: 75 }, (_, idx) => `antigravity:11:s-${idx + 175}`), + ...Array.from({ length: 3 }, (_, idx) => `antigravity:11:g-${idx}`), + ], + ); + // Page statements are untraced, so the trace does not grow with the tables. + assert.isTrue(spans.some((span) => span.name === "sqlite")); + assert.deepStrictEqual( + spans.flatMap((span) => { + const query = span.attributes.get("db.query.text"); + return typeof query === "string" && query.includes("rowid > ?") ? [query] : []; + }), + [], + ); + // Every row, with metadata or not, comes back through a bounded statement. + for (const [table, total, most] of [ + ["steps", 250, 64], + ["gen_metadata", 3, 1], + ] as const) { + const rows = reads + .filter((read) => read.sql.includes(`FROM ${table} `)) + .map((read) => read.rows); + assert.strictEqual( + rows.reduce((sum, count) => sum + count, 0), + total, + table, + ); + assert.isAtMost(Math.max(...rows), most, table); + } + }).pipe(Effect.provide(NodeServices.layer)), + ); + + it.effect("returns Antigravity records in idx order however the rows are stored", () => + Effect.gen(function* () { + const path = yield* Path.Path; + const dir = yield* tempDir; + yield* seedDatabase(path.join(dir, "session-1.db"), (db) => + Effect.gen(function* () { + yield* createTables(db); + yield* db.exec("BEGIN"); + // Stored in reverse idx order, so rowid order is the opposite of idx order. + for (let idx = 99; idx >= 0; idx--) { + yield* db.run("INSERT INTO steps VALUES (?, ?)", [idx, stepUsage(`s-${idx}`)]); + } + for (const idx of [1, 0]) { + yield* db.run("INSERT INTO gen_metadata VALUES (?, ?)", [ + idx, + antigravityGeneration(`g-${idx}`), + ]); + } + yield* db.exec("COMMIT"); + }), + ); + const result = yield* readAntigravityUsage(dir, 0); + assert.deepStrictEqual(result.errors, []); + assert.deepStrictEqual( + records(result).map((record) => record.dedupeKey), + [ + ...Array.from({ length: 100 }, (_, idx) => `antigravity:11:s-${idx}`), + "antigravity:11:g-0", + "antigravity:11:g-1", + ], + ); + }).pipe(Effect.provide(NodeServices.layer)), + ); }); diff --git a/apps/server/src/provider/Drivers/antigravityUsageReader.ts b/apps/server/src/provider/Drivers/antigravityUsageReader.ts index e48609cf874d23416ef8f9c03dfe4285f05ba967..5239d5d551c33ee09ab08319688e2c2658cc7a01 100644 --- a/apps/server/src/provider/Drivers/antigravityUsageReader.ts +++ b/apps/server/src/provider/Drivers/antigravityUsageReader.ts @@ -132,11 +132,12 @@ function metadata(bytes: Uint8Array, step: boolean): 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)]; + // Copy the usage bytes: subarrays kept in usages would hold each row's whole blob alive. + const usages = usage === undefined ? [] : [fields(usage.slice())]; 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)); + if (retryUsage !== undefined) usages.push(fields(retryUsage.slice())); } return { model: modelName( @@ -171,6 +172,14 @@ class AntigravityDatabaseDecodeError extends Schema.TaggedError = []; - for (const row of yield* sql.unsafe>(query)) { - const idx = row.idx; - if (typeof idx !== "number") { - return yield* new AntigravityDatabaseDecodeError({ path }); + // Below every rowid. + let after = -Infinity; + while (true) { + // Untraced: a cold read runs one of these per generation. + const rows = yield* sql + .unsafe<{ + readonly row_id: number; + readonly idx: unknown; + readonly data: unknown; + }>(query, [after, pageRows]) + .pipe(Effect.withTracerEnabled(false)); + for (const row of rows) { + // Skipped here rather than in SQL, so a statement never scans past + // more than one page of steps that have no metadata. + if (step && row.data === null) continue; + const idx = row.idx; + if (typeof idx !== "number") { + return yield* new AntigravityDatabaseDecodeError({ path }); + } + entries.push({ idx, entry: yield* decode(() => metadata(blob(row.data), step)) }); } - entries.push({ idx, entry: yield* decode(() => metadata(blob(row[column]), step)) }); - if (entries.length % 256 === 0) yield* Effect.yieldNow; + if (rows.length < pageRows) break; + after = rows[rows.length - 1]!.row_id; + yield* Effect.yieldNow; } - return entries; + return entries.sort((a, b) => a.idx - b.idx); }); const generations = tables.has("gen_metadata") - ? yield* readMetadata("SELECT idx, data FROM gen_metadata ORDER BY idx", "data", false) + ? yield* readMetadata( + "SELECT rowid AS row_id, idx, data FROM gen_metadata WHERE rowid > ? ORDER BY rowid LIMIT ?", + GENERATION_PAGE_ROWS, + false, + ) : []; let trajectoryTimestamp: number | null = null; if (tables.has("trajectory_metadata_blob")) { @@ -221,8 +255,8 @@ const readDatabase = Effect.fn("readDatabase")(function* ( } const steps = tables.has("steps") ? yield* readMetadata( - "SELECT idx, metadata FROM steps WHERE metadata IS NOT NULL ORDER BY idx", - "metadata", + "SELECT rowid AS row_id, idx, metadata AS data FROM steps WHERE rowid > ? ORDER BY rowid LIMIT ?", + STEP_PAGE_ROWS, true, ) : []; diff --git a/packages/provider-opencode/src/server/usage.test.ts b/packages/provider-opencode/src/server/usage.test.ts index 4dc7ec99486c6d5bb41bc4259fa5f7cd7377d725..e2e80e42712c7f518bac22ae53592e79434d53c9 100644 --- a/packages/provider-opencode/src/server/usage.test.ts +++ b/packages/provider-opencode/src/server/usage.test.ts @@ -87,6 +87,62 @@ describe("readOpenCodeUsage", () => { }).pipe(Effect.provide(NodeServices.layer)), ); + it.effect("reads at most 64 OpenCode rowids per statement however much history is older", () => + Effect.gen(function* () { + const fileSystem = yield* FileSystem.FileSystem; + const path = yield* Path.Path; + const dir = yield* fileSystem.makeTempDirectoryScoped({ prefix: "usage-reader-test-" }); + const db = yield* openDatabase(path.join(dir, "opencode.db")); + yield* db.unsafe( + "CREATE TABLE message (id TEXT, session_id TEXT, data TEXT, time_created INTEGER)", + ); + const sinceMs = 1780000000000; + const data = JSON.stringify({ + role: "assistant", + modelID: "claude-sonnet-4-5", + tokens: { input: 100, output: 20 }, + }); + // OpenCode appends, so the history before the window comes first. + yield* db.unsafe("BEGIN"); + for (let index = 0; index < 305; index++) { + yield* db.unsafe("INSERT INTO message VALUES (?, ?, ?, ?)", [ + `msg-${index}`, + "session-1", + data, + sinceMs + index - 300, + ]); + } + yield* db.unsafe("COMMIT"); + + const spans: number[] = []; + const originalAll = NodeSqlite.StatementSync.prototype.all; + yield* Effect.acquireRelease( + Effect.sync(() => + vi.spyOn(NodeSqlite.StatementSync.prototype, "all").mockImplementation(function ( + this: NodeSqlite.StatementSync, + ...params + ) { + const rows = originalAll.apply(this, params); + if (this.sourceSQL.includes("FROM message WHERE")) { + const range = /rowid > (-?[\d.]+) AND rowid <= (-?[\d.]+)/.exec(this.expandedSQL); + spans.push(range ? Number(range[2]) - Number(range[1]) : Infinity); + } + return rows; + }), + ), + (spy) => Effect.sync(() => spy.mockRestore()), + ); + const result = yield* readOpenCodeUsage(dir, sinceMs); + assert.isFalse(result.error); + assert.deepStrictEqual( + result.files.flatMap((file) => file.records).map((record) => record.dedupeKey), + Array.from({ length: 5 }, (_, index) => `opencode:msg-${index + 300}`), + ); + assert.isNotEmpty(spans); + assert.isAtMost(Math.max(...spans), 64); + }).pipe(Effect.provide(NodeServices.layer)), + ); + it.effect("counts migrated OpenCode messages once and sees subsequent WAL writes", () => Effect.gen(function* () { const fileSystem = yield* FileSystem.FileSystem; diff --git a/packages/provider-opencode/src/server/usage.ts b/packages/provider-opencode/src/server/usage.ts index 63b5568f05285f88bbab622705c0987fc9dda09c..6ccfb688891e082cf94fadef2fc5bb37eda53990 100644 --- a/packages/provider-opencode/src/server/usage.ts +++ b/packages/provider-opencode/src/server/usage.ts @@ -188,34 +188,37 @@ export const readOpenCodeUsage = Effect.fn("readOpenCodeUsage")(function* ( const timestamp = columns.has("time_created") ? "time_created" : "NULL"; const predicates = table === "session_message" ? ["type = 'assistant'"] : []; if (timestamp !== "NULL") predicates.push("time_created >= ?"); - const lastRowId = (yield* sql.unsafe<{ readonly last_row_id: number | null }>( - `SELECT MAX(rowid) AS last_row_id FROM ${table}`, - ))[0]?.last_row_id; - if (lastRowId === null || lastRowId === undefined) continue; - let afterRowId: number | undefined; - // SqlClient materializes each result with StatementSync.all(). Keep - // message bodies in small pages, and exclude rows appended mid-scan. - while (true) { - const pagePredicates = [ - ...predicates, - "rowid <= ?", - ...(afterRowId === undefined ? [] : ["rowid > ?"]), - ]; - const rows = yield* sql.unsafe<{ - readonly row_id: number; - readonly id: unknown; - readonly session_id: unknown; - readonly data: unknown; - readonly created: unknown; - }>( - `SELECT rowid AS row_id, id, session_id, data, ${timestamp} AS created - FROM ${table} WHERE ${pagePredicates.join(" AND ")} ORDER BY rowid LIMIT ${MESSAGE_PAGE_SIZE}`, - [ - ...(timestamp === "NULL" ? [] : [sinceMs]), - lastRowId, - ...(afterRowId === undefined ? [] : [afterRowId]), - ], - ); + predicates.push("rowid > ?", "rowid <= ?"); + const [range] = yield* sql.unsafe<{ + readonly first_row_id: number | null; + readonly last_row_id: number | null; + }>( + `SELECT (SELECT MIN(rowid) FROM ${table}) AS first_row_id, (SELECT MAX(rowid) FROM ${table}) AS last_row_id`, + ); + const firstRowId = range?.first_row_id; + const lastRowId = range?.last_row_id; + if (typeof firstRowId !== "number" || typeof lastRowId !== "number") continue; + // SqlClient materializes each result with StatementSync.all(), and the + // window can exclude most rows. Each statement covers MESSAGE_PAGE_SIZE + // rowids, however few match, and rows appended mid-scan are left out. + // Untraced, as a large history runs thousands of them. + for (let after = firstRowId - 1; after < lastRowId; after += MESSAGE_PAGE_SIZE) { + const rows = yield* sql + .unsafe<{ + readonly id: unknown; + readonly session_id: unknown; + readonly data: unknown; + readonly created: unknown; + }>( + `SELECT id, session_id, data, ${timestamp} AS created + FROM ${table} WHERE ${predicates.join(" AND ")} ORDER BY rowid`, + [ + ...(timestamp === "NULL" ? [] : [sinceMs]), + after, + Math.min(after + MESSAGE_PAGE_SIZE, lastRowId), + ], + ) + .pipe(Effect.withTracerEnabled(false)); for (const row of rows) { append( file.records, @@ -226,8 +229,6 @@ export const readOpenCodeUsage = Effect.fn("readOpenCodeUsage")(function* ( }), ); } - if (rows.length < MESSAGE_PAGE_SIZE) break; - afterRowId = rows[rows.length - 1]!.row_id; yield* Effect.yieldNow; } } -- 2.54.0