authorgravatar for me@t3.ggTheo Browne <me@t3.gg> 2026-10-14 15:53:24-07:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2026-10-14 15:53:24-07:00
log52ffe0cae6e5e03c95037152ef4548f40daf5c2e
treed79d48986ca8b51c0b24ced93c4441355674c87f
parentf4df79477709185c490627a3fd43af96674ae778
signature Signed by PGP key B5690EEEBB952194

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 <noreply@anthropic.com>

4 files changed, 274 insertions(+), 44 deletions(-)

apps/server/src/provider/Drivers/antigravityUsageReader.test.ts+140-1
...@@ -1,8 +1,9 @@...@@ -1,8 +1,9 @@
1// @effect-diagnostics-next-line nodeBuiltinImport:off - Effect's File.Info has no ctime or sub-millisecond mtime, which the cache tests compare.1// @effect-diagnostics-next-line nodeBuiltinImport:off - Effect's File.Info has no ctime or sub-millisecond mtime, which the cache tests compare.
2import * as NodeFSP from "node:fs/promises";2import * as NodeFSP from "node:fs/promises";
3import * as NodeSqlite from "node:sqlite";
34
4import * as NodeServices from "@effect/platform-node/NodeServices";5import * as NodeServices from "@effect/platform-node/NodeServices";
5import { assert, describe, it } from "@effect/vitest";6import { assert, describe, it, vi } from "@effect/vitest";
6import * as NodeSqliteClient from "@t3tools/shared/nodeSqliteClient";7import * as NodeSqliteClient from "@t3tools/shared/nodeSqliteClient";
7import * as Context from "effect/Context";8import * as Context from "effect/Context";
8import * as Effect from "effect/Effect";9import * as Effect from "effect/Effect";
...@@ -10,6 +11,7 @@ import * as FileSystem from "effect/FileSystem";...@@ -10,6 +11,7 @@ import * as FileSystem from "effect/FileSystem";
10import * as Layer from "effect/Layer";11import * as Layer from "effect/Layer";
11import * as Path from "effect/Path";12import * as Path from "effect/Path";
12import * as SqlClient from "effect/sql/SqlClient";13import * as SqlClient from "effect/sql/SqlClient";
14import * as Tracer from "effect/Tracer";
1315
14import { makeAntigravityUsageCache, readAntigravityUsage } from "./antigravityUsageReader.ts";16import { makeAntigravityUsageCache, readAntigravityUsage } from "./antigravityUsageReader.ts";
1517
...@@ -445,4 +447,141 @@ describe("readAntigravityUsage", () => {...@@ -445,4 +447,141 @@ describe("readAntigravityUsage", () => {
445 assert.strictEqual(records(result)[0]?.totals.uncachedInputTokens, 10);447 assert.strictEqual(records(result)[0]?.totals.uncachedInputTokens, 10);
446 }).pipe(Effect.provide(NodeServices.layer)),448 }).pipe(Effect.provide(NodeServices.layer)),
447 );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 );
448});587});
apps/server/src/provider/Drivers/antigravityUsageReader.ts+47-13
...@@ -132,11 +132,12 @@ function metadata(bytes: Uint8Array, step: boolean): Metadata {...@@ -132,11 +132,12 @@ function metadata(bytes: Uint8Array, step: boolean): Metadata {
132 const data = step ? root : nested(root, 1);132 const data = step ? root : nested(root, 1);
133 const model = step ? nested(data, 24) : data;133 const model = step ? nested(data, 24) : data;
134 const usage = bytesAt(data, step ? 9 : 4);134 const usage = bytesAt(data, step ? 9 : 4);
135 const usages = usage === undefined ? [] : [fields(usage)];135 // Copy the usage bytes: subarrays kept in usages would hold each row's whole blob alive.
136 const usages = usage === undefined ? [] : [fields(usage.slice())];
136 for (const retry of data.get(step ? 28 : 17) ?? []) {137 for (const retry of data.get(step ? 28 : 17) ?? []) {
137 if (!(retry instanceof Uint8Array)) throw new Error("Invalid retry metadata");138 if (!(retry instanceof Uint8Array)) throw new Error("Invalid retry metadata");
138 const retryUsage = bytesAt(fields(retry), 2);139 const retryUsage = bytesAt(fields(retry), 2);
139 if (retryUsage !== undefined) usages.push(fields(retryUsage));140 if (retryUsage !== undefined) usages.push(fields(retryUsage.slice()));
140 }141 }
141 return {142 return {
142 model: modelName(143 model: modelName(
...@@ -171,6 +172,14 @@ class AntigravityDatabaseDecodeError extends Schema.TaggedError<AntigravityDatab...@@ -171,6 +172,14 @@ class AntigravityDatabaseDecodeError extends Schema.TaggedError<AntigravityDatab
171 }172 }
172}173}
173174
175/**
176 * Rows per statement. One statement.all() holds the event loop until every row
177 * it returns is in memory, and a generation row carries the conversation's
178 * whole context so far, often megabytes, so generations are read one at a time.
179 */
180const GENERATION_PAGE_ROWS = 1;
181const STEP_PAGE_ROWS = 64;
182
174const readDatabase = Effect.fn("readDatabase")(function* (183const readDatabase = Effect.fn("readDatabase")(function* (
175 path: string,184 path: string,
176 sessionId: string,185 sessionId: string,
...@@ -192,24 +201,49 @@ const readDatabase = Effect.fn("readDatabase")(function* (...@@ -192,24 +201,49 @@ const readDatabase = Effect.fn("readDatabase")(function* (
192 if (!tables.has("gen_metadata") && !tables.has("steps")) {201 if (!tables.has("gen_metadata") && !tables.has("steps")) {
193 return yield* new AntigravityDatabaseDecodeError({ path });202 return yield* new AntigravityDatabaseDecodeError({ path });
194 }203 }
204 // Runs `query` a page at a time, binding the last rowid read and the page
205 // size. A rowid seek works whatever the table's indexes. Antigravity declares
206 // idx as the INTEGER PRIMARY KEY, so rowid order is idx order; the sort keeps
207 // idx order where rows are stored otherwise.
195 const readMetadata = Effect.fnUntraced(function* (208 const readMetadata = Effect.fnUntraced(function* (
196 query: string,209 query: string,
197 column: string,210 pageRows: number,
198 step: boolean,211 step: boolean,
199 ) {212 ) {
200 const entries: Array<{ idx: number; entry: Metadata }> = [];213 const entries: Array<{ idx: number; entry: Metadata }> = [];
201 for (const row of yield* sql.unsafe<Record<string, unknown>>(query)) {214 // Below every rowid.
202 const idx = row.idx;215 let after = -Infinity;
203 if (typeof idx !== "number") {216 while (true) {
204 return yield* new AntigravityDatabaseDecodeError({ path });217 // Untraced: a cold read runs one of these per generation.
218 const rows = yield* sql
219 .unsafe<{
220 readonly row_id: number;
221 readonly idx: unknown;
222 readonly data: unknown;
223 }>(query, [after, pageRows])
224 .pipe(Effect.withTracerEnabled(false));
225 for (const row of rows) {
226 // Skipped here rather than in SQL, so a statement never scans past
227 // more than one page of steps that have no metadata.
228 if (step && row.data === null) continue;
229 const idx = row.idx;
230 if (typeof idx !== "number") {
231 return yield* new AntigravityDatabaseDecodeError({ path });
232 }
233 entries.push({ idx, entry: yield* decode(() => metadata(blob(row.data), step)) });
205 }234 }
206 entries.push({ idx, entry: yield* decode(() => metadata(blob(row[column]), step)) });235 if (rows.length < pageRows) break;
207 if (entries.length % 256 === 0) yield* Effect.yieldNow;236 after = rows[rows.length - 1]!.row_id;
237 yield* Effect.yieldNow;
208 }238 }
209 return entries;239 return entries.sort((a, b) => a.idx - b.idx);
210 });240 });
211 const generations = tables.has("gen_metadata")241 const generations = tables.has("gen_metadata")
212 ? yield* readMetadata("SELECT idx, data FROM gen_metadata ORDER BY idx", "data", false)242 ? yield* readMetadata(
243 "SELECT rowid AS row_id, idx, data FROM gen_metadata WHERE rowid > ? ORDER BY rowid LIMIT ?",
244 GENERATION_PAGE_ROWS,
245 false,
246 )
213 : [];247 : [];
214 let trajectoryTimestamp: number | null = null;248 let trajectoryTimestamp: number | null = null;
215 if (tables.has("trajectory_metadata_blob")) {249 if (tables.has("trajectory_metadata_blob")) {
...@@ -221,8 +255,8 @@ const readDatabase = Effect.fn("readDatabase")(function* (...@@ -221,8 +255,8 @@ const readDatabase = Effect.fn("readDatabase")(function* (
221 }255 }
222 const steps = tables.has("steps")256 const steps = tables.has("steps")
223 ? yield* readMetadata(257 ? yield* readMetadata(
224 "SELECT idx, metadata FROM steps WHERE metadata IS NOT NULL ORDER BY idx",258 "SELECT rowid AS row_id, idx, metadata AS data FROM steps WHERE rowid > ? ORDER BY rowid LIMIT ?",
225 "metadata",259 STEP_PAGE_ROWS,
226 true,260 true,
227 )261 )
228 : [];262 : [];
packages/provider-opencode/src/server/usage.test.ts+56
...@@ -87,6 +87,62 @@ describe("readOpenCodeUsage", () => {...@@ -87,6 +87,62 @@ describe("readOpenCodeUsage", () => {
87 }).pipe(Effect.provide(NodeServices.layer)),87 }).pipe(Effect.provide(NodeServices.layer)),
88 );88 );
8989
90 it.effect("reads at most 64 OpenCode rowids per statement however much history is older", () =>
91 Effect.gen(function* () {
92 const fileSystem = yield* FileSystem.FileSystem;
93 const path = yield* Path.Path;
94 const dir = yield* fileSystem.makeTempDirectoryScoped({ prefix: "usage-reader-test-" });
95 const db = yield* openDatabase(path.join(dir, "opencode.db"));
96 yield* db.unsafe(
97 "CREATE TABLE message (id TEXT, session_id TEXT, data TEXT, time_created INTEGER)",
98 );
99 const sinceMs = 1780000000000;
100 const data = JSON.stringify({
101 role: "assistant",
102 modelID: "claude-sonnet-4-5",
103 tokens: { input: 100, output: 20 },
104 });
105 // OpenCode appends, so the history before the window comes first.
106 yield* db.unsafe("BEGIN");
107 for (let index = 0; index < 305; index++) {
108 yield* db.unsafe("INSERT INTO message VALUES (?, ?, ?, ?)", [
109 `msg-${index}`,
110 "session-1",
111 data,
112 sinceMs + index - 300,
113 ]);
114 }
115 yield* db.unsafe("COMMIT");
116
117 const spans: number[] = [];
118 const originalAll = NodeSqlite.StatementSync.prototype.all;
119 yield* Effect.acquireRelease(
120 Effect.sync(() =>
121 vi.spyOn(NodeSqlite.StatementSync.prototype, "all").mockImplementation(function (
122 this: NodeSqlite.StatementSync,
123 ...params
124 ) {
125 const rows = originalAll.apply(this, params);
126 if (this.sourceSQL.includes("FROM message WHERE")) {
127 const range = /rowid > (-?[\d.]+) AND rowid <= (-?[\d.]+)/.exec(this.expandedSQL);
128 spans.push(range ? Number(range[2]) - Number(range[1]) : Infinity);
129 }
130 return rows;
131 }),
132 ),
133 (spy) => Effect.sync(() => spy.mockRestore()),
134 );
135 const result = yield* readOpenCodeUsage(dir, sinceMs);
136 assert.isFalse(result.error);
137 assert.deepStrictEqual(
138 result.files.flatMap((file) => file.records).map((record) => record.dedupeKey),
139 Array.from({ length: 5 }, (_, index) => `opencode:msg-${index + 300}`),
140 );
141 assert.isNotEmpty(spans);
142 assert.isAtMost(Math.max(...spans), 64);
143 }).pipe(Effect.provide(NodeServices.layer)),
144 );
145
90 it.effect("counts migrated OpenCode messages once and sees subsequent WAL writes", () =>146 it.effect("counts migrated OpenCode messages once and sees subsequent WAL writes", () =>
91 Effect.gen(function* () {147 Effect.gen(function* () {
92 const fileSystem = yield* FileSystem.FileSystem;148 const fileSystem = yield* FileSystem.FileSystem;
packages/provider-opencode/src/server/usage.ts+31-30
...@@ -188,34 +188,37 @@ export const readOpenCodeUsage = Effect.fn("readOpenCodeUsage")(function* (...@@ -188,34 +188,37 @@ export const readOpenCodeUsage = Effect.fn("readOpenCodeUsage")(function* (
188 const timestamp = columns.has("time_created") ? "time_created" : "NULL";188 const timestamp = columns.has("time_created") ? "time_created" : "NULL";
189 const predicates = table === "session_message" ? ["type = 'assistant'"] : [];189 const predicates = table === "session_message" ? ["type = 'assistant'"] : [];
190 if (timestamp !== "NULL") predicates.push("time_created >= ?");190 if (timestamp !== "NULL") predicates.push("time_created >= ?");
191 const lastRowId = (yield* sql.unsafe<{ readonly last_row_id: number | null }>(191 predicates.push("rowid > ?", "rowid <= ?");
192 `SELECT MAX(rowid) AS last_row_id FROM ${table}`,192 const [range] = yield* sql.unsafe<{
193 ))[0]?.last_row_id;193 readonly first_row_id: number | null;
194 if (lastRowId === null || lastRowId === undefined) continue;194 readonly last_row_id: number | null;
195 let afterRowId: number | undefined;195 }>(
196 // SqlClient materializes each result with StatementSync.all(). Keep196 `SELECT (SELECT MIN(rowid) FROM ${table}) AS first_row_id, (SELECT MAX(rowid) FROM ${table}) AS last_row_id`,
197 // message bodies in small pages, and exclude rows appended mid-scan.197 );
198 while (true) {198 const firstRowId = range?.first_row_id;
199 const pagePredicates = [199 const lastRowId = range?.last_row_id;
200 ...predicates,200 if (typeof firstRowId !== "number" || typeof lastRowId !== "number") continue;
201 "rowid <= ?",201 // SqlClient materializes each result with StatementSync.all(), and the
202 ...(afterRowId === undefined ? [] : ["rowid > ?"]),202 // window can exclude most rows. Each statement covers MESSAGE_PAGE_SIZE
203 ];203 // rowids, however few match, and rows appended mid-scan are left out.
204 const rows = yield* sql.unsafe<{204 // Untraced, as a large history runs thousands of them.
205 readonly row_id: number;205 for (let after = firstRowId - 1; after < lastRowId; after += MESSAGE_PAGE_SIZE) {
206 readonly id: unknown;206 const rows = yield* sql
207 readonly session_id: unknown;207 .unsafe<{
208 readonly data: unknown;208 readonly id: unknown;
209 readonly created: unknown;209 readonly session_id: unknown;
210 }>(210 readonly data: unknown;
211 `SELECT rowid AS row_id, id, session_id, data, ${timestamp} AS created211 readonly created: unknown;
212 FROM ${table} WHERE ${pagePredicates.join(" AND ")} ORDER BY rowid LIMIT ${MESSAGE_PAGE_SIZE}`,212 }>(
213 [213 `SELECT id, session_id, data, ${timestamp} AS created
214 ...(timestamp === "NULL" ? [] : [sinceMs]),214 FROM ${table} WHERE ${predicates.join(" AND ")} ORDER BY rowid`,
215 lastRowId,215 [
216 ...(afterRowId === undefined ? [] : [afterRowId]),216 ...(timestamp === "NULL" ? [] : [sinceMs]),
217 ],217 after,
218 );218 Math.min(after + MESSAGE_PAGE_SIZE, lastRowId),
219 ],
220 )
221 .pipe(Effect.withTracerEnabled(false));
219 for (const row of rows) {222 for (const row of rows) {
220 append(223 append(
221 file.records,224 file.records,
...@@ -226,8 +229,6 @@ export const readOpenCodeUsage = Effect.fn("readOpenCodeUsage")(function* (...@@ -226,8 +229,6 @@ export const readOpenCodeUsage = Effect.fn("readOpenCodeUsage")(function* (
226 }),229 }),
227 );230 );
228 }231 }
229 if (rows.length < MESSAGE_PAGE_SIZE) break;
230 afterRowId = rows[rows.length - 1]!.row_id;
231 yield* Effect.yieldNow;232 yield* Effect.yieldNow;
232 }233 }
233 }234 }