diff --git a/packages/provider-opencode/src/server/usage.test.ts b/packages/provider-opencode/src/server/usage.test.ts index 20a768609a7da..6139bc615a006 100644 --- a/packages/provider-opencode/src/server/usage.test.ts +++ b/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; diff --git a/packages/provider-opencode/src/server/usage.ts b/packages/provider-opencode/src/server/usage.ts index be7e23a10a8a8..07f65d1b6871e 100644 --- a/packages/provider-opencode/src/server/usage.ts +++ b/packages/provider-opencode/src/server/usage.ts @@ -111,6 +111,8 @@ export const readOpenCodeUsage = Effect.fn("readOpenCodeUsage")(function* ( const seen = new Set(); 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,19 +172,20 @@ 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; @@ -190,13 +193,15 @@ export const readOpenCodeUsage = Effect.fn("readOpenCodeUsage")(function* ( 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 }; });