perf(opencode): bound usage reads to in-window messages

packages/provider-opencode/src/server/usage.test.ts

@@ -25,6 +25,7 @@ describe("readOpenCodeUsage", () => {
       const path = yield* Path.Path;
       const dir = yield* fileSystem.makeTempDirectoryScoped({ prefix: "usage-reader-large-test-" });
       const db = yield* openDatabase(path.join(dir, "opencode.db"));
+      yield* db.unsafe("PRAGMA journal_mode = WAL");
       yield* db.unsafe(
         "CREATE TABLE message (id TEXT, session_id TEXT, data TEXT, time_created INTEGER)",
       );
@@ -34,7 +35,7 @@ describe("readOpenCodeUsage", () => {
         modelID: "claude-sonnet-4-5",
         time: { created: sinceMs },
         tokens: { input: 100, output: 20 },
-        body: "x".repeat(64 * 1024),
+        body: "x".repeat(256 * 1024),
       };
       const data = JSON.stringify(message);
       const insert = (rowId: number, id: string, payload: string, created = sinceMs) =>
@@ -47,9 +48,24 @@ describe("readOpenCodeUsage", () => {
       }
       yield* insert(300, "old", data, sinceMs - 1);
       yield* insert(301, "user", JSON.stringify({ ...message, role: "user" }));
+      for (let index = 0; index < 32; index++) {
+        yield* insert(
+          400 + index,
+          `user-${index}`,
+          JSON.stringify({
+            role: "user",
+            summary: { diffs: [{ patch: "x".repeat(1024 * 1024) }] },
+          }),
+        );
+      }
       yield* insert(302, "invalid", "{invalid");
       yield* insert(303, "msg-0", data);
 
+      const writer = yield* Effect.acquireRelease(
+        Effect.sync(() => new NodeSqlite.DatabaseSync(path.join(dir, "opencode.db"))),
+        (database) => Effect.sync(() => database.close()),
+      );
+      let changed = false;
       let largestRawResult = 0;
       const originalAll = NodeSqlite.StatementSync.prototype.all;
       yield* Effect.acquireRelease(
@@ -64,6 +80,12 @@ describe("readOpenCodeUsage", () => {
               0,
             );
             largestRawResult = Math.max(largestRawResult, rawSize);
+            if (!changed && rawSize > 0) {
+              changed = true;
+              writer
+                .prepare("UPDATE message SET data = ? WHERE id = ?")
+                .run(JSON.stringify({ ...message, tokens: { input: 900, output: 20 } }), "msg-129");
+            }
             return rows;
           }),
         ),
@@ -80,8 +102,65 @@ describe("readOpenCodeUsage", () => {
         records.reduce((total, record) => total + record.totals.uncachedInputTokens, 0),
         13000,
       );
+      assert.strictEqual(
+        (yield* readOpenCodeUsage(dir, sinceMs)).files
+          .flatMap((file) => file.records)
+          .find((record) => record.dedupeKey === "opencode:msg-129")?.totals.uncachedInputTokens,
+        900,
+      );
       assert.isAbove(largestRawResult, 0);
-      assert.isBelow(largestRawResult, 6 * 1024 * 1024);
+      assert.isBelow(largestRawResult, 64 * 1024);
+    }).pipe(Effect.provide(NodeServices.layer)),
+  );
+
+  it.effect("reads statements in proportion to the window despite sparse, unordered rowids", () =>
+    Effect.gen(function* () {
+      const fileSystem = yield* FileSystem.FileSystem;
+      const path = yield* Path.Path;
+      const dir = yield* fileSystem.makeTempDirectoryScoped({ prefix: "usage-window-" });
+      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)",
+      );
+      yield* db.unsafe("CREATE INDEX message_time ON message(time_created)");
+      const sinceMs = 1780000000000;
+      const data = JSON.stringify({ modelID: "m", tokens: { input: 1 } });
+      yield* db.unsafe(
+        `WITH RECURSIVE n(i) AS (SELECT 1 UNION ALL SELECT i + 1 FROM n WHERE i < 5000)
+        INSERT INTO message SELECT 'old-' || i, 's', ?, ? FROM n`,
+        [data, sinceMs - 1],
+      );
+      for (const [rowid, id] of [
+        [-10, "a"],
+        [10000, "b"],
+        [20000, "c"],
+      ] as const) {
+        yield* db.unsafe(
+          "INSERT INTO message(rowid, id, session_id, data, time_created) VALUES (?, ?, 's', ?, ?)",
+          [rowid, id, data, sinceMs],
+        );
+      }
+      let statements = 0;
+      const originalAll = NodeSqlite.StatementSync.prototype.all;
+      yield* Effect.acquireRelease(
+        Effect.sync(() =>
+          vi.spyOn(NodeSqlite.StatementSync.prototype, "all").mockImplementation(function (
+            this: NodeSqlite.StatementSync,
+            ...params
+          ) {
+            statements++;
+            return originalAll.apply(this, params);
+          }),
+        ),
+        (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),
+        ["opencode:a", "opencode:b", "opencode:c"],
+      );
+      assert.isBelow(statements, 10);
     }).pipe(Effect.provide(NodeServices.layer)),
   );
 
@@ -136,6 +215,61 @@ describe("readOpenCodeUsage", () => {
     }).pipe(Effect.provide(NodeServices.layer)),
   );
 
+  it.effect(
+    "projects both SQLite formats without changing malformed or optional field handling",
+    () =>
+      Effect.gen(function* () {
+        const fileSystem = yield* FileSystem.FileSystem;
+        const path = yield* Path.Path;
+        const dir = yield* fileSystem.makeTempDirectoryScoped({ prefix: "usage-formats-" });
+        const db = yield* openDatabase(path.join(dir, "opencode.db"));
+        const sinceMs = 1780000000000;
+        const message = { model: { id: "m" }, tokens: { input: 10 }, time: { created: sinceMs } };
+        const cases = [
+          ["assistant", JSON.stringify({ ...message, role: "assistant", cost: 0.25 }), true],
+          ["absent-role", JSON.stringify(message), true],
+          ["null-role", JSON.stringify({ ...message, role: null }), false],
+          ["user", JSON.stringify({ ...message, role: "user" }), false],
+          ["malformed", "{broken", false],
+          ["json5", "{role:'assistant'}", false],
+          ["null", null, false],
+          ["blob", new Uint8Array(Buffer.from(JSON.stringify(message))), false],
+          ["fallback", JSON.stringify({ modelID: "fallback-model", tokens: { output: 5 } }), true],
+        ] as const;
+        for (const table of ["message", "session_message"]) {
+          yield* db.unsafe(
+            `CREATE TABLE ${table} (id TEXT, session_id TEXT, data, type TEXT, time_created INTEGER)`,
+          );
+          for (const [id, data] of cases) {
+            yield* db.unsafe(`INSERT INTO ${table} VALUES (?, 's', ?, 'assistant', ?)`, [
+              `${table}-${id}`,
+              data,
+              sinceMs,
+            ]);
+          }
+          yield* db.unsafe(`INSERT INTO ${table} VALUES (?, 's', ?, 'user', ?)`, [
+            `${table}-type-user`,
+            JSON.stringify(message),
+            sinceMs,
+          ]);
+        }
+        const result = yield* readOpenCodeUsage(dir, sinceMs);
+        assert.isFalse(result.error);
+        const records = result.files.flatMap((file) => file.records);
+        assert.deepStrictEqual(
+          records.map((record) => record.dedupeKey),
+          [
+            ...cases.filter(([, , valid]) => valid).map(([id]) => `opencode:message-${id}`),
+            "opencode:message-type-user",
+            ...cases.filter(([, , valid]) => valid).map(([id]) => `opencode:session_message-${id}`),
+          ],
+        );
+        assert.strictEqual(records[0]?.reportedCostUsd, 0.25);
+        assert.strictEqual(records[2]?.model, "fallback-model");
+        assert.strictEqual(records[2]?.timestampMs, sinceMs);
+      }).pipe(Effect.provide(NodeServices.layer)),
+  );
+
   it.effect("reads only legacy OpenCode messages written since the window opened", () =>
     Effect.gen(function* () {
       const fileSystem = yield* FileSystem.FileSystem;

packages/provider-opencode/src/server/usage.ts

@@ -111,6 +111,8 @@ export const readOpenCodeUsage = Effect.fn("readOpenCodeUsage")(function* (
   const seen = new Set<string>();
   let found = false;
   let error = false;
+  let databaseRows = 0;
+  let databasePages = 0;
   const append = (records: UsageRecord[], record: UsageRecord | null) => {
     if (record === null || record.timestampMs < sinceMs) return;
     if (record.dedupeKey !== null) {
@@ -170,33 +172,36 @@ 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 >= ?");
-        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;
-        // StatementSync.all materializes results; page by rowid and exclude mid-scan appends.
-        // Omit tracing because large histories require thousands of statements.
-        for (let after = firstRowId - 1; after < lastRowId; after += MESSAGE_PAGE_SIZE) {
+        // List only rowids so SQLite can use time indexes without reading message bodies.
+        const listed = (yield* sql.unsafe<{ readonly row_id: number }>(
+          `SELECT rowid AS row_id FROM ${table}
+           ${predicates.length ? `WHERE ${predicates.join(" AND ")}` : ""}`,
+          timestamp === "NULL" ? [] : [sinceMs],
+        )).toSorted((a, b) => a.row_id - b.row_id);
+        databaseRows += listed.length;
+        // Filter and project in SQLite: user diffs and assistant bodies can be hundreds of MB.
+        const validAssistant = `CASE WHEN typeof(data) = 'text' AND json_valid(data)
+          THEN json_type(data, '$.role') IS NULL OR data ->> '$.role' = 'assistant' END`;
+        // Omit per-page spans; the enclosing scan records row and page counts.
+        for (let start = 0; start < listed.length; start += MESSAGE_PAGE_SIZE) {
+          const page = listed.slice(start, start + MESSAGE_PAGE_SIZE);
+          databasePages++;
           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),
-              ],
+              `SELECT id, session_id, ${timestamp} AS created, json_object(
+                 'tokens', data -> '$.tokens', 'cost', data -> '$.cost',
+                 'model', data -> '$.model', 'modelID', data -> '$.modelID',
+                 'time', data -> '$.time', 'id', data -> '$.id', 'sessionID', data -> '$.sessionID'
+               ) AS data
+               FROM ${table}
+               WHERE rowid IN (${page.map(() => "?").join(", ")}) AND ${validAssistant}
+               ORDER BY rowid`,
+              page.map((row) => row.row_id),
             )
             .pipe(Effect.withTracerEnabled(false));
           for (const row of rows) {
@@ -213,6 +218,7 @@ export const readOpenCodeUsage = Effect.fn("readOpenCodeUsage")(function* (
         }
       }
     }).pipe(
+      (read) => Effect.flatMap(SqlClient.SqlClient, (sql) => sql.withTransaction(read)),
       Effect.provide(NodeSqliteClient.layer({ filename: file.path, readonly: true })),
       Effect.catchTags({
         SqlError: () => {
@@ -288,6 +294,10 @@ export const readOpenCodeUsage = Effect.fn("readOpenCodeUsage")(function* (
       }),
     );
   }
+  yield* Effect.annotateCurrentSpan({
+    "usage.database_rows": databaseRows,
+    "usage.database_pages": databasePages,
+  });
   return { files, missing: !found && !error, error };
 });