diff --git a/packages/provider-opencode/src/server/usage.test.ts b/packages/provider-opencode/src/server/usage.test.ts index d211563738ad5..02d1e5dc8cc28 100644 --- a/packages/provider-opencode/src/server/usage.test.ts +++ b/packages/provider-opencode/src/server/usage.test.ts @@ -1,3 +1,4 @@ +import type * as NodeWorkerThreads from "node:worker_threads"; import * as NodeServices from "@effect/platform-node/NodeServices"; import { assert, describe, it } from "@effect/vitest"; import { DEFAULT_SERVER_SETTINGS, ProviderInstanceId } from "@t3tools/contracts"; @@ -13,6 +14,27 @@ import { vi } from "vite-plus/test"; import { openCodeUsageReader, readOpenCodeUsage } from "./usage.ts"; +const readerWorkers = vi.hoisted(() => ({ started: 0, enabled: false })); + +// Statement spies only see the main thread, so reader workers fail to start unless a test enables them. +vi.mock("node:worker_threads", async (importOriginal) => { + const actual = await importOriginal(); + class ReaderWorker extends actual.Worker { + constructor(...args: ConstructorParameters) { + super( + ...((readerWorkers.enabled + ? args + : [ + "throw new Error('reader workers disabled');", + { eval: true }, + ]) as ConstructorParameters), + ); + readerWorkers.started++; + } + } + return { ...actual, Worker: ReaderWorker }; +}); + const openDatabase = (filename: string) => Layer.build(NodeSqliteClient.layer({ filename })).pipe( Effect.map(Context.get(SqlClient.SqlClient)), @@ -141,6 +163,7 @@ describe("readOpenCodeUsage", () => { ); } let statements = 0; + let bodyReads = 0; const originalAll = NodeSqlite.StatementSync.prototype.all; yield* Effect.acquireRelease( Effect.sync(() => @@ -149,6 +172,7 @@ describe("readOpenCodeUsage", () => { ...params ) { statements++; + if (this.sourceSQL.includes("data ->")) bodyReads++; return originalAll.apply(this, params); }), ), @@ -161,6 +185,51 @@ describe("readOpenCodeUsage", () => { ["opencode:a", "opencode:b", "opencode:c"], ); assert.isBelow(statements, 10); + + bodyReads = 0; + const repeated = yield* readOpenCodeUsage(dir, sinceMs); + assert.deepStrictEqual( + repeated.files.flatMap((file) => file.records).map((record) => record.dedupeKey), + ["opencode:a", "opencode:b", "opencode:c"], + ); + // An unchanged database must not have its message bodies read and parsed again. + assert.strictEqual(bodyReads, 0); + }).pipe(Effect.provide(NodeServices.layer)), + ); + + it.effect("does not reuse rows from a replaced or restored database file", () => + Effect.gen(function* () { + const fileSystem = yield* FileSystem.FileSystem; + const path = yield* Path.Path; + const dir = yield* fileSystem.makeTempDirectoryScoped({ prefix: "usage-replaced-" }); + const write = (name: string, input: number) => + Effect.scoped( + Effect.gen(function* () { + const db = yield* openDatabase(path.join(dir, name)); + yield* db.unsafe("PRAGMA journal_mode = WAL"); + yield* db.unsafe( + "CREATE TABLE message (id TEXT, session_id TEXT, time_created INTEGER, time_updated INTEGER, data TEXT)", + ); + yield* db.unsafe( + "INSERT INTO message(rowid, id, session_id, time_created, time_updated, data) VALUES (1, 'a', 's', 1, 1, ?)", + [JSON.stringify({ modelID: "m", tokens: { input } })], + ); + }), + ); + const input = Effect.map( + readOpenCodeUsage(dir, 0), + (result) => result.files[0]?.records[0]?.totals.uncachedInputTokens, + ); + yield* write("opencode.db", 100); + assert.strictEqual(yield* input, 100); + yield* write("replacement.db", 900); + yield* fileSystem.rename(path.join(dir, "replacement.db"), path.join(dir, "opencode.db")); + assert.strictEqual(yield* input, 900); + // A backup restored in place keeps the inode and birth time, like `cp -p`. + yield* write("backup.db", 500); + yield* fileSystem.copyFile(path.join(dir, "backup.db"), path.join(dir, "opencode.db")); + yield* fileSystem.utimes(path.join(dir, "opencode.db"), 0, 0); + assert.strictEqual(yield* input, 500); }).pipe(Effect.provide(NodeServices.layer)), ); @@ -188,6 +257,27 @@ describe("readOpenCodeUsage", () => { const legacy = path.join(dir, "storage", "message", message.sessionID); yield* fileSystem.makeDirectory(legacy, { recursive: true }); yield* fileSystem.writeFileString(path.join(legacy, "msg-1.json"), JSON.stringify(message)); + yield* Effect.acquireRelease( + Effect.sync(() => { + readerWorkers.enabled = true; + readerWorkers.started = 0; + }), + () => Effect.sync(() => (readerWorkers.enabled = false)), + ); + let mainThreadBodyReads = 0; + const originalAll = NodeSqlite.StatementSync.prototype.all; + yield* Effect.acquireRelease( + Effect.sync(() => + vi.spyOn(NodeSqlite.StatementSync.prototype, "all").mockImplementation(function ( + this: NodeSqlite.StatementSync, + ...params + ) { + if (this.sourceSQL.includes("data ->")) mainThreadBodyReads++; + return originalAll.apply(this, params); + }), + ), + (spy) => Effect.sync(() => spy.mockRestore()), + ); const first = yield* readOpenCodeUsage(dir, 0); assert.isFalse(first.error); const records = first.files.flatMap((file) => file.records); @@ -212,6 +302,9 @@ describe("readOpenCodeUsage", () => { ["opencode:msg-2"], ); assert.isAbove(Number((yield* fileSystem.stat(path.join(dir, "opencode.db-wal"))).size), 0); + // A large database must not be read on the server's event loop. + assert.isAbove(readerWorkers.started, 0); + assert.strictEqual(mainThreadBodyReads, 0); }).pipe(Effect.provide(NodeServices.layer)), ); diff --git a/packages/provider-opencode/src/server/usage.ts b/packages/provider-opencode/src/server/usage.ts index 4ae7b6df358e9..8d9bb28c70b34 100644 --- a/packages/provider-opencode/src/server/usage.ts +++ b/packages/provider-opencode/src/server/usage.ts @@ -11,6 +11,7 @@ import * as FileSystem from "effect/FileSystem"; import * as Option from "effect/Option"; import * as Path from "effect/Path"; import type * as PlatformError from "effect/PlatformError"; +import * as Semaphore from "effect/Semaphore"; import * as SqlClient from "effect/sql/SqlClient"; import type { OpenCodeSettings } from "../settings.ts"; @@ -93,6 +94,18 @@ export interface OpenCodeUsageReadResult { readonly error: boolean; } +interface CachedRow { + readonly updated: unknown; + readonly size: unknown; + readonly created: unknown; + readonly record: UsageRecord | null; +} + +// Parsed rows by `${databasePath}\0${table}`, then rowid, so unchanged rows are not read again. +const rowCache = new Map }>(); +// Concurrent scans of the same database would each read every uncached row. +const scanLock = Semaphore.makeUnsafe(1); + const isNotFound = (cause: PlatformError.PlatformError) => cause.reason._tag === "NotFound"; const MESSAGE_PAGE_SIZE = 64; @@ -116,6 +129,7 @@ const LEGACY_STAT_BATCH = 256; export const readOpenCodeUsage = Effect.fn("readOpenCodeUsage")(function* ( root: string, sinceMs: number, + retentionCutoffMs = sinceMs, ): Effect.fn.Return { const fileSystem = yield* FileSystem.FileSystem; const path = yield* Path.Path; @@ -164,6 +178,22 @@ export const readOpenCodeUsage = Effect.fn("readOpenCodeUsage")(function* ( found = true; const file = { path: path.join(root, name), records: [] as UsageRecord[] }; files.push(file); + // Taken before the snapshot opens, so a write that races the scan invalidates its rows. + const [main, wal] = yield* Effect.forEach([file.path, `${file.path}-wal`], (target) => + Effect.option(fileSystem.stat(target)), + ); + const stamp = (info: Option.Option) => + Option.match(info, { + onNone: () => "none", + onSome: (stat) => `${stat.size}:${Option.getOrNull(stat.mtime)?.getTime()}`, + }); + // A replaced or restored database can repeat rowids and change markers, so every key includes + // the main file; WAL commits leave it alone until a checkpoint. + const identity = Option.match(main!, { + onNone: () => "none", + onSome: (stat) => + `${stat.dev}:${Option.getOrNull(stat.ino)}:${Option.getOrNull(stat.birthtime)?.getTime()}/${stamp(main!)}`, + }); yield* Effect.gen(function* () { const sql = yield* SqlClient.SqlClient; // Fail promptly on live-writer locks instead of stalling the server. @@ -182,30 +212,53 @@ export const readOpenCodeUsage = Effect.fn("readOpenCodeUsage")(function* ( ), ); const timestamp = columns.has("time_created") ? "time_created" : "NULL"; + const updated = columns.has("time_updated") ? "time_updated" : "NULL"; + // Without time_updated an in-place edit can keep the length, so any database write counts. + const version = updated === "NULL" ? `${identity}/${stamp(wal!)}` : identity; const predicates = table === "session_message" ? ["type = 'assistant'"] : []; if (timestamp !== "NULL") predicates.push("time_created >= ?"); - // 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} + // List rowids and change markers so SQLite can use time indexes without reading message + // bodies; octet_length comes from the record header. + const listed = (yield* sql.unsafe<{ + readonly row_id: number; + readonly updated: unknown; + readonly size: unknown; + readonly created: unknown; + }>( + `SELECT rowid AS row_id, ${updated} AS updated, octet_length(data) AS size, + ${timestamp} AS created + FROM ${table} ${predicates.length ? `WHERE ${predicates.join(" AND ")}` : ""}`, timestamp === "NULL" ? [] : [sinceMs], )).toSorted((a, b) => a.row_id - b.row_id); databaseRows += listed.length; + const key = `${file.path}\0${table}`; + const previous = rowCache.get(key); + const reusable = previous?.version === version ? previous.rows : undefined; + const cached = new Map(); + const misses: typeof listed = []; + for (const row of listed) { + const hit = reusable?.get(row.row_id); + if (hit && hit.updated === row.updated && hit.size === row.size) + cached.set(row.row_id, hit); + else misses.push(row); + } // 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); + for (let start = 0; start < misses.length; start += MESSAGE_PAGE_SIZE) { + const page = misses.slice(start, start + MESSAGE_PAGE_SIZE); databasePages++; const rows = yield* sql .unsafe<{ + readonly row_id: number; readonly id: unknown; readonly session_id: unknown; readonly data: unknown; readonly created: unknown; }>( - `SELECT id, session_id, ${timestamp} AS created, json_object( + `SELECT rowid AS row_id, id, session_id, ${timestamp} AS created, json_object( 'tokens', data -> '$.tokens', 'cost', data -> '$.cost', 'model', data -> '$.model', 'modelID', data -> '$.modelID', 'providerID', data -> '$.providerID', @@ -217,9 +270,10 @@ export const readOpenCodeUsage = Effect.fn("readOpenCodeUsage")(function* ( page.map((row) => row.row_id), ) .pipe(Effect.withTracerEnabled(false)); + const parsed = new Map(); for (const row of rows) { - append( - file.records, + parsed.set( + row.row_id, parseOpenCodeMessage(text(row.data), { id: text(row.id), sessionId: text(row.session_id), @@ -227,12 +281,28 @@ export const readOpenCodeUsage = Effect.fn("readOpenCodeUsage")(function* ( }), ); } + // Rows the SQL filter dropped are cached as empty so they are not read again either. + for (const row of page) { + cached.set(row.row_id, { ...row, record: parsed.get(row.row_id) ?? null }); + } yield* Effect.yieldNow; } + for (const row of listed) append(file.records, cached.get(row.row_id)!.record); + // Rows before this window were not listed; keep them for a wider window until retention. + for (const [rowId, entry] of reusable ?? []) { + const created = entry.created; + if (typeof created === "number" && created < sinceMs && created >= retentionCutoffMs) + cached.set(rowId, entry); + } + rowCache.set(key, { version, rows: cached }); } }).pipe( (read) => Effect.flatMap(SqlClient.SqlClient, (sql) => sql.withTransaction(read)), - Effect.provide(NodeSqliteClient.layer({ filename: file.path, readonly: true })), + // Run the scan on a reader worker so large databases do not block the server's event loop. + NodeSqliteClient.readOnly, + Effect.provide( + NodeSqliteClient.layer({ filename: file.path, readonly: true, readerWorker: true }), + ), Effect.catchTags({ SqlError: () => { error = true; @@ -312,7 +382,7 @@ export const readOpenCodeUsage = Effect.fn("readOpenCodeUsage")(function* ( "usage.database_pages": databasePages, }); return { files, missing: !found && !error, error }; -}); +}, scanLock.withPermit); const resolveOpenCodeDataDirs = Effect.fn("resolveOpenCodeDataDirs")(function* ( environment: NodeJS.ProcessEnv, @@ -346,7 +416,11 @@ export type OpenCodeUsageReaderEnv = FileSystem.FileSystem | Path.Path; export const openCodeUsageReader: ProviderUsageReader = { kind: "scan", provider: "opencode", - scan: Effect.fn("openCodeUsageReader.scan")(function* ({ instances, windowStartMs }) { + scan: Effect.fn("openCodeUsageReader.scan")(function* ({ + instances, + windowStartMs, + retentionCutoffMs, + }) { const roots = new Set(); for (const instance of instances) { for (const root of yield* resolveOpenCodeDataDirs(instance.environment)) roots.add(root); @@ -354,7 +428,7 @@ export const openCodeUsageReader: ProviderUsageReader - readOpenCodeUsage(dir, windowStartMs).pipe( + readOpenCodeUsage(dir, windowStartMs, retentionCutoffMs).pipe( Effect.map((result) => ({ dir, files: result.missing && !result.error ? null : result.files,