tokenmax-collector 0.2.0 → 0.2.1

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "tokenmax-collector",
-  "version": "0.2.0",
+  "version": "0.2.1",
   "description": "Report local ccusage token usage to a tokenmax instance",
   "type": "module",
   "bin": {
@@ -19,7 +19,7 @@
     "directory": "packages/collector"
   },
   "engines": {
-    "bun": ">=1.3.14"
+    "bun": ">=1.4.0"
   },
   "scripts": {
     "check-types": "tsc --noEmit",

package/src/antigravity.ts

@@ -11,6 +11,11 @@ export interface AntigravityStep {
   output: number;
 }
 
+export interface AntigravityUsage {
+  failures: string[];
+  steps: AntigravityStep[];
+}
+
 interface ProtoField {
   number: number;
   value: bigint | Uint8Array;
@@ -43,10 +48,16 @@ export function antigravityConversationsDir(home: string): string {
 function decodeMessage(bytes: Uint8Array): ProtoField[] {
   const fields: ProtoField[] = [];
   let offset = 0;
+  const truncated = (): never => {
+    throw new Error("truncated protobuf message");
+  };
   const varint = (): bigint => {
     let result = 0n;
     let shift = 0n;
     for (;;) {
+      if (offset >= bytes.length) {
+        truncated();
+      }
       const byte = bytes[offset++];
       result |= BigInt(byte & 0x7f) << shift;
       if (byte < 0x80) {
@@ -56,6 +67,9 @@ function decodeMessage(bytes: Uint8Array): ProtoField[] {
     }
   };
   const bytesOf = (length: number): Uint8Array => {
+    if (offset + length > bytes.length) {
+      truncated();
+    }
     const slice = bytes.subarray(offset, offset + length);
     offset += length;
     return slice;
@@ -108,7 +122,7 @@ function text(fields: ProtoField[], number: number): string | null {
 function openReadOnly(file: string): DatabaseSync {
   const location = existsSync(`${file}-wal`)
     ? file
-    : `file:${file}?immutable=1`;
+    : `file:${file.split("/").map(encodeURIComponent).join("/")}?immutable=1`;
   return new DatabaseSync(location, { readOnly: true });
 }
 
@@ -159,23 +173,35 @@ function readConversation(
 
 export async function readAntigravitySteps(
   dir: string,
-): Promise<AntigravityStep[]> {
+): Promise<AntigravityUsage> {
   let files: string[];
   try {
     files = await readdir(dir);
   } catch (error) {
     if ((error as NodeJS.ErrnoException).code === "ENOENT") {
-      return [];
+      return { failures: [], steps: [] };
     }
     throw error;
   }
   const names = new Map<number, string>();
   const steps: RawStep[] = [];
+  const failures: string[] = [];
   for (const file of files.filter((name) => name.endsWith(".db")).sort()) {
-    readConversation(resolve(dir, file), names, steps);
+    const conversation = resolve(dir, file);
+    const before = steps.length;
+    try {
+      readConversation(conversation, names, steps);
+    } catch (error) {
+      steps.length = before;
+      const message = error instanceof Error ? error.message : String(error);
+      failures.push(`${conversation}: ${message}`);
+    }
   }
-  return steps.map(({ modelCode, ...step }) => ({
-    ...step,
-    model: names.get(modelCode) ?? unknownModel,
-  }));
+  return {
+    failures,
+    steps: steps.map(({ modelCode, ...step }) => ({
+      ...step,
+      model: names.get(modelCode) ?? unknownModel,
+    })),
+  };
 }

package/src/cli.ts

@@ -36,12 +36,15 @@ async function runCollect(io: CliIo): Promise<number> {
     runner: io.runner,
     today: io.today,
   });
-  if (result.kind === "reported") {
-    io.stdout(`accepted ${result.accepted} days for ${result.machine}`);
-    return 0;
-  }
-  if (result.kind === "empty") {
-    io.stdout("nothing to report");
+  if (result.kind === "reported" || result.kind === "empty") {
+    for (const warning of result.warnings) {
+      io.stderr(warning);
+    }
+    io.stdout(
+      result.kind === "reported"
+        ? `accepted ${result.accepted} days for ${result.machine}`
+        : "nothing to report",
+    );
     return 0;
   }
   if (result.kind === "missing-config") {

package/src/collect.ts

@@ -27,8 +27,8 @@ export interface CollectOptions {
 }
 
 export type CollectResult =
-  | { kind: "reported"; accepted: number; machine: string }
-  | { kind: "empty" }
+  | { kind: "reported"; accepted: number; machine: string; warnings: string[] }
+  | { kind: "empty"; warnings: string[] }
   | { kind: "missing-config"; configFile: string }
   | { kind: "failed"; message: string };
 
@@ -56,11 +56,17 @@ function acceptedCount(body: string): number | null {
   return typeof accepted === "number" ? accepted : null;
 }
 
+interface AntigravityRows {
+  days: UsageDay[];
+  warnings: string[];
+}
+
 async function report(
   fetcher: Fetcher,
   url: string,
   key: string,
   usage: UsageReport,
+  warnings: string[],
 ): Promise<CollectResult> {
   const response = await fetcher(reportUrl(url), {
     body: JSON.stringify(usage),
@@ -84,7 +90,7 @@ async function report(
       message: `tokenmax responded an unexpected body: ${body}`,
     };
   }
-  return { accepted, kind: "reported", machine: usage.machine };
+  return { accepted, kind: "reported", machine: usage.machine, warnings };
 }
 
 async function antigravityDays(
@@ -93,16 +99,20 @@ async function antigravityDays(
   timezone: string,
   fetcher: Fetcher,
   pricesFile: string,
-): Promise<UsageDay[]> {
+): Promise<AntigravityRows> {
   const since = windowStart(today, timezone);
-  const steps = (
-    await readAntigravitySteps(antigravityConversationsDir(home))
-  ).filter((step) => calendarDate(step.at, timezone) >= since);
+  const usage = await readAntigravitySteps(antigravityConversationsDir(home));
+  const warnings = usage.failures.map(
+    (failure) => `antigravity: skipped ${failure}`,
+  );
+  const steps = usage.steps.filter(
+    (step) => calendarDate(step.at, timezone) >= since,
+  );
   if (steps.length === 0) {
-    return [];
+    return { days: [], warnings };
   }
   const prices = await loadPrices(fetcher, pricesFile, today);
-  return mapAntigravitySteps(steps, timezone, prices);
+  return { days: mapAntigravitySteps(steps, timezone, prices), warnings };
 }
 
 export async function collect(options: CollectOptions): Promise<CollectResult> {
@@ -133,29 +143,36 @@ export async function collect(options: CollectOptions): Promise<CollectResult> {
       sinceArgument(today, timezone),
       timezone,
     );
-    days = [
-      ...mapCcusageDays(daily),
-      ...(await antigravityDays(
-        env.home,
-        today,
-        timezone,
-        fetcher,
-        paths.pricesFile,
-      )),
-    ];
+    days = mapCcusageDays(daily);
   } catch (error) {
     return { kind: "failed", message: messageOf(error) };
   }
+
+  let antigravity: AntigravityRows;
+  try {
+    antigravity = await antigravityDays(
+      env.home,
+      today,
+      timezone,
+      fetcher,
+      paths.pricesFile,
+    );
+  } catch (error) {
+    antigravity = { days: [], warnings: [`antigravity: ${messageOf(error)}`] };
+  }
+  days.push(...antigravity.days);
   if (days.length === 0) {
-    return { kind: "empty" };
+    return { kind: "empty", warnings: antigravity.warnings };
   }
 
   try {
-    return await report(fetcher, config.config.url, config.config.key, {
-      days,
-      machine,
-      timezone,
-    });
+    return await report(
+      fetcher,
+      config.config.url,
+      config.config.key,
+      { days, machine, timezone },
+      antigravity.warnings,
+    );
   } catch (error) {
     return { kind: "failed", message: messageOf(error) };
   }

package/src/pricing.ts

@@ -1,4 +1,4 @@
-import { mkdir, readFile, stat, writeFile } from "node:fs/promises";
+import { mkdir, readFile, rename, stat, writeFile } from "node:fs/promises";
 import { dirname } from "node:path";
 import { z } from "zod";
 import type { Fetcher } from "./http";
@@ -57,19 +57,32 @@ export function costOf(price: ModelPrice, usage: PricedUsage): number {
 
 async function readCache(
   pricesFile: string,
-): Promise<{ modifiedAt: number; source: string } | null> {
+): Promise<{ modifiedAt: number; table: PriceTable } | null> {
+  let source: string;
+  let modifiedAt: number;
   try {
-    const [source, info] = await Promise.all([
+    [source, modifiedAt] = await Promise.all([
       readFile(pricesFile, "utf8"),
-      stat(pricesFile),
+      stat(pricesFile).then((info) => info.mtimeMs),
     ]);
-    return { modifiedAt: info.mtimeMs, source };
   } catch (error) {
     if ((error as NodeJS.ErrnoException).code === "ENOENT") {
       return null;
     }
     throw error;
   }
+  try {
+    return { modifiedAt, table: parseLitellmPrices(source) };
+  } catch {
+    return null;
+  }
+}
+
+async function writeCache(pricesFile: string, source: string): Promise<void> {
+  await mkdir(dirname(pricesFile), { recursive: true });
+  const partial = `${pricesFile}.${process.pid}.tmp`;
+  await writeFile(partial, source);
+  await rename(partial, pricesFile);
 }
 
 async function fetchPrices(fetcher: Fetcher): Promise<string> {
@@ -88,7 +101,7 @@ export async function loadPrices(
 ): Promise<PriceTable> {
   const cached = await readCache(pricesFile);
   if (cached !== null && now.getTime() - cached.modifiedAt < maxAgeMs) {
-    return parseLitellmPrices(cached.source);
+    return cached.table;
   }
   let source: string;
   let table: PriceTable;
@@ -97,12 +110,11 @@ export async function loadPrices(
     table = parseLitellmPrices(source);
   } catch (error) {
     if (cached !== null) {
-      return parseLitellmPrices(cached.source);
+      return cached.table;
     }
     const message = error instanceof Error ? error.message : String(error);
     throw new Error(`could not load model prices: ${message}`);
   }
-  await mkdir(dirname(pricesFile), { recursive: true });
-  await writeFile(pricesFile, source);
+  await writeCache(pricesFile, source);
   return table;
 }