perf(server): Antigravity usage skips re-decoding unchanged databases after restart

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

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 = <A, E>(
+        use: (
+          scan: () => Effect.Effect<
+            Effect.Success<ReturnType<AntigravityUsage.AntigravityUsage["Service"]["scan"]>>
+          >,
+        ) => Effect.Effect<A, E>,
+      ) =>
+        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<any>(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;

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<string, CachedDatabase>();
 
+// 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<string, CachedDatabase>) =>
+  encodeCacheFile({
+    models: MODEL_IDS_SIGNATURE,
+    databases: Object.fromEntries(databases),
+  });
+
 export const readAntigravityUsage = Effect.fn("readAntigravityUsage")(function* (
   conversationsDirectories: string | readonly string[],
   sinceMs: number,