diff --git a/apps/server/src/provider/Drivers/AntigravityUsage.ts b/apps/server/src/provider/Drivers/AntigravityUsage.ts index 02168fdd83149..5925a9de8417b 100644 --- a/apps/server/src/provider/Drivers/AntigravityUsage.ts +++ b/apps/server/src/provider/Drivers/AntigravityUsage.ts @@ -6,6 +6,7 @@ import type { ProviderUsageReader, ProviderUsageScan, } from "@t3tools/provider-core/server/usage"; +import { writeFileStringAtomically } from "@t3tools/shared/atomicWrite"; import * as HostProcess from "@t3tools/shared/HostProcess"; import * as Context from "effect/Context"; import * as Crypto from "effect/Crypto"; @@ -15,7 +16,13 @@ import * as Layer from "effect/Layer"; import * as Path from "effect/Path"; import { resolveAntigravityInstanceDirectories } from "../antigravityAuthSupport.ts"; -import { makeAntigravityUsageCache, readAntigravityUsage } from "./antigravityUsageReader.ts"; +import { + ANTIGRAVITY_USAGE_CACHE_FILE_NAME, + decodeAntigravityUsageCache, + encodeAntigravityUsageCache, + makeAntigravityUsageCache, + readAntigravityUsage, +} from "./antigravityUsageReader.ts"; export class AntigravityUsage extends Context.Service< AntigravityUsage, @@ -33,6 +40,20 @@ const make = Effect.gen(function* () { const path = yield* Path.Path; const host = yield* ProviderHost.ProviderHost; const cache = makeAntigravityUsageCache(); + const cachePath = path.join(host.paths.stateDir, ANTIGRAVITY_USAGE_CACHE_FILE_NAME); + // What the cache file holds; a failed write leaves it behind so the next scan retries. + let persisted = new Map(cache); + // Restore decoded databases once, so a restart re-reads only changed ones. + const loadCache = yield* Effect.cached( + fileSystem.readFileString(cachePath).pipe( + Effect.flatMap(decodeAntigravityUsageCache), + Effect.map(({ databases }) => { + for (const [key, entry] of Object.entries(databases)) cache.set(key, entry); + persisted = new Map(cache); + }), + Effect.ignore, + ), + ); const dataRoots = Effect.fn("AntigravityUsage.dataRoots")(function* ( environment: NodeJS.ProcessEnv, @@ -93,10 +114,28 @@ const make = Effect.gen(function* () { : resolvedRoot; conversationDirs.add(yield* fileSystem.realPath(dir).pipe(Effect.orElseSucceed(() => dir))); } + yield* loadCache; const result = yield* readAntigravityUsage([...conversationDirs], windowStartMs, cache).pipe( Effect.provideService(FileSystem.FileSystem, fileSystem), Effect.provideService(Path.Path, path), ); + if ( + cache.size !== persisted.size || + cache.entries().some(([key, entry]) => persisted.get(key) !== entry) + ) { + const snapshot = new Map(cache); + yield* encodeAntigravityUsageCache(snapshot).pipe( + Effect.flatMap((contents) => writeFileStringAtomically({ filePath: cachePath, contents })), + Effect.tap(() => + Effect.sync(() => { + persisted = snapshot; + }), + ), + Effect.provideService(FileSystem.FileSystem, fileSystem), + Effect.provideService(Path.Path, path), + Effect.ignore, + ); + } const scanned: ProviderUsageScan[] = []; for (const dir of conversationDirs) { const exists = yield* fileSystem diff --git a/apps/server/src/provider/Drivers/antigravityUsageReader.test.ts b/apps/server/src/provider/Drivers/antigravityUsageReader.test.ts index a08eac31746ff..59074d29ee21f 100644 --- a/apps/server/src/provider/Drivers/antigravityUsageReader.test.ts +++ b/apps/server/src/provider/Drivers/antigravityUsageReader.test.ts @@ -3,15 +3,29 @@ import * as NodeFSP from "node:fs/promises"; import * as NodeServices from "@effect/platform-node/NodeServices"; import { assert, describe, it } from "@effect/vitest"; +import { ProviderInstanceId } from "@t3tools/contracts"; import * as NodeSqliteClient from "@t3tools/shared/nodeSqliteClient"; import * as Context from "effect/Context"; import * as Effect from "effect/Effect"; +import * as Exit from "effect/Exit"; 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 { makeAntigravityUsageCache, readAntigravityUsage } from "./antigravityUsageReader.ts"; +import * as ServerSecretStore from "../../auth/ServerSecretStore.ts"; +import * as BackgroundPolicy from "../../background/BackgroundPolicy.ts"; +import * as ServerConfig from "../../config.ts"; +import * as ServerSettings from "../../serverSettings.ts"; +import * as ProviderHostLive from "../ProviderHostLive.ts"; +import * as AntigravityUsage from "./AntigravityUsage.ts"; +import { + ANTIGRAVITY_USAGE_CACHE_FILE_NAME, + decodeAntigravityUsageCache, + encodeAntigravityUsageCache, + makeAntigravityUsageCache, + readAntigravityUsage, +} from "./antigravityUsageReader.ts"; const openDatabase = Effect.fn("openDatabase")(function* (filename: string) { const context = yield* Layer.build(NodeSqliteClient.layer({ filename })); @@ -339,6 +353,94 @@ describe("readAntigravityUsage", () => { ]), ); + it.effect("persists the service cache, retrying after a failed write, and restores it", () => + Effect.gen(function* () { + const fileSystem = yield* FileSystem.FileSystem; + const path = yield* Path.Path; + const root = yield* tempDir; + const baseDir = path.join(root, "state"); + const configLayer = ServerConfig.layerTest(root, baseDir); + const { stateDir } = yield* ServerConfig.ServerConfig.pipe(Effect.provide(configLayer)); + const conversations = path.join(root, "data", "conversations"); + yield* fileSystem.makeDirectory(conversations, { recursive: true }); + yield* seedDatabase(path.join(conversations, "session-1.db"), (db) => + db + .exec("CREATE TABLE gen_metadata (idx INTEGER, data BLOB)") + .pipe( + Effect.andThen( + db.run("INSERT INTO gen_metadata VALUES (0, ?)", [antigravityGeneration("r-1")]), + ), + ), + ); + // Open the database once so the scans below see a stable fingerprint. + yield* readAntigravityUsage(conversations, 0); + const cacheFile = path.join(stateDir, ANTIGRAVITY_USAGE_CACHE_FILE_NAME); + const instances = [ + { + instanceId: ProviderInstanceId.make("antigravity"), + config: undefined, + environment: { ANTIGRAVITY_DATA_DIR: path.join(root, "data") }, + configured: false, + }, + ]; + let failures = 1; + const flakyFileSystem = Layer.effect( + FileSystem.FileSystem, + Effect.map(FileSystem.FileSystem, (inner) => ({ + ...inner, + // Writing under a missing directory fails with a real PlatformError. + writeFileString: (target: string, contents: string) => + failures-- > 0 + ? inner.writeFileString(path.join(target, "missing", "file"), contents) + : inner.writeFileString(target, contents), + })), + ); + const layer = AntigravityUsage.layer.pipe( + Layer.provideMerge(flakyFileSystem), + Layer.provideMerge(ProviderHostLive.layer), + Layer.provideMerge(Layer.mock(BackgroundPolicy.BackgroundPolicy)({})), + Layer.provideMerge(Layer.mock(ServerSecretStore.ServerSecretStore)({})), + Layer.provideMerge(configLayer), + Layer.provideMerge(ServerSettings.layerTest({})), + Layer.provideMerge(NodeServices.layer), + ); + const withService = ( + use: ( + scan: () => Effect.Effect< + Effect.Success> + >, + ) => Effect.Effect, + ) => + Effect.gen(function* () { + const service = yield* AntigravityUsage.AntigravityUsage; + return yield* use(() => service.scan({ instances, windowStartMs: 0 }).pipe(Effect.orDie)); + }).pipe(Effect.provide(layer)); + // One service: the first write fails, and an unchanged second scan must write again. + const first = yield* withService((scan) => + Effect.gen(function* () { + const result = yield* scan(); + assert.isFalse(yield* fileSystem.exists(cacheFile)); + yield* scan(); + return result; + }), + ); + const saved = yield* fileSystem.readFileString(cacheFile); + const databases = (yield* decodeAntigravityUsageCache(saved)).databases; + assert.strictEqual(Object.keys(databases).length, 1); + // A new service on the same state dir must take its records from the saved file, not a rescan: + // an edited cached record shows up in the result while the database fingerprint is unchanged. + const edited = JSON.parse(saved); + for (const entry of Object.values(edited.databases)) { + for (const candidate of entry.candidates) + candidate.record.sessionId = "restored-from-cache"; + } + yield* fileSystem.writeFileString(cacheFile, JSON.stringify(edited)); + const restored = yield* withService((scan) => scan()); + assert.notInclude(JSON.stringify(first), "restored-from-cache"); + assert.include(JSON.stringify(restored), "restored-from-cache"); + }).pipe(Effect.provide(NodeServices.layer)), + ); + it.effect("rereads an Antigravity database when only its WAL changed", () => Effect.gen(function* () { const pathService = yield* Path.Path; @@ -369,6 +471,38 @@ describe("readAntigravityUsage", () => { }).pipe(Effect.provide(NodeServices.layer)), ); + it.effect("restores a persisted cache without rereading unchanged databases", () => + Effect.gen(function* () { + const dir = yield* tempDir; + const path = (yield* Path.Path).join(dir, "session-1.db"); + yield* seedDatabase(path, (db) => + db + .exec("CREATE TABLE gen_metadata (idx INTEGER, data BLOB)") + .pipe( + Effect.andThen( + db.run("INSERT INTO gen_metadata VALUES (0, ?)", [antigravityGeneration("r-1")]), + ), + ), + ); + const cache = makeAntigravityUsageCache(); + const first = yield* readAntigravityUsage(dir, 0, cache); + const saved = yield* encodeAntigravityUsageCache(cache); + const restored = makeAntigravityUsageCache(); + for (const [key, entry] of Object.entries( + (yield* decodeAntigravityUsageCache(saved)).databases, + )) { + restored.set(key, entry); + } + const [entry] = restored.values(); + assert.deepStrictEqual(yield* readAntigravityUsage(dir, 0, restored), first); + assert.strictEqual([...restored.values()][0], entry); + // Unreadable files, and files written for another MODEL_IDS table, are not used. + for (const bad of ["{", saved.replace(/"models":"[^"]*"/, '"models":""')]) { + assert.isTrue(Exit.isFailure(yield* Effect.exit(decodeAntigravityUsageCache(bad)))); + } + }).pipe(Effect.provide(NodeServices.layer)), + ); + it.effect("reads Antigravity step-only stores and reports malformed databases", () => Effect.gen(function* () { const fileSystem = yield* FileSystem.FileSystem; diff --git a/apps/server/src/provider/Drivers/antigravityUsageReader.ts b/apps/server/src/provider/Drivers/antigravityUsageReader.ts index d3035d5ae9b5f..5d4da76551346 100644 --- a/apps/server/src/provider/Drivers/antigravityUsageReader.ts +++ b/apps/server/src/provider/Drivers/antigravityUsageReader.ts @@ -2,6 +2,7 @@ import * as NodeFSP from "node:fs/promises"; import { pipe } from "effect/Function"; +import { UsageProviderKind, UsageTokenTotals } from "@t3tools/contracts"; import type { UsageRecord } from "@t3tools/provider-core/server/usage"; import * as NodeSqliteClient from "@t3tools/shared/nodeSqliteClient"; import * as Effect from "effect/Effect"; @@ -332,6 +333,47 @@ interface CachedDatabase { /** Fingerprint canonical DB/WAL paths by size, mtime and ctime; WAL holds new rows, and ctime catches restored mtimes. */ export const makeAntigravityUsageCache = () => new Map(); +// A UsageRecord without its optional fields, as persisted in the cache file. +const CachedUsageRecord = Schema.Struct({ + provider: UsageProviderKind, + timestampMs: Schema.Number, + model: Schema.String, + sessionId: Schema.String, + totals: UsageTokenTotals, + reportedCostUsd: Schema.NullOr(Schema.Number), + speed: Schema.Literals(["standard", "fast", "ultrafast"]), + dedupeKey: Schema.NullOr(Schema.String), +}); +// Rename when usageCandidates output changes; cached models also pin MODEL_IDS. +export const ANTIGRAVITY_USAGE_CACHE_FILE_NAME = "usage-antigravity-cache-v1.json"; +const MODEL_IDS_SIGNATURE = Object.entries(MODEL_IDS).join(); +// Plain records and arrays; ReadonlyMap codecs made encode and decode several times slower. +const AntigravityUsageCacheFile = Schema.fromJsonString( + Schema.Struct({ + models: Schema.Literal(MODEL_IDS_SIGNATURE), + databases: Schema.Record( + Schema.String, + Schema.Struct({ + fingerprint: Schema.String, + candidates: Schema.Array( + Schema.Struct({ + record: CachedUsageRecord, + keys: Schema.Array(Schema.String), + timestampQuality: Schema.Number, + }), + ), + }), + ), + }), +); +export const decodeAntigravityUsageCache = Schema.decodeUnknownEffect(AntigravityUsageCacheFile); +const encodeCacheFile = Schema.encodeEffect(AntigravityUsageCacheFile); +export const encodeAntigravityUsageCache = (databases: ReadonlyMap) => + encodeCacheFile({ + models: MODEL_IDS_SIGNATURE, + databases: Object.fromEntries(databases), + }); + export const readAntigravityUsage = Effect.fn("readAntigravityUsage")(function* ( conversationsDirectories: string | readonly string[], sinceMs: number,