tokenmax-collector 0.1.0 → 0.2.1

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "tokenmax-collector",
-  "version": "0.1.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

@@ -0,0 +1,207 @@
+import { existsSync } from "node:fs";
+import { readdir } from "node:fs/promises";
+import { resolve } from "node:path";
+import { DatabaseSync } from "node:sqlite";
+
+export interface AntigravityStep {
+  at: Date;
+  cacheRead: number;
+  input: number;
+  model: string;
+  output: number;
+}
+
+export interface AntigravityUsage {
+  failures: string[];
+  steps: AntigravityStep[];
+}
+
+interface ProtoField {
+  number: number;
+  value: bigint | Uint8Array;
+}
+
+interface RawStep {
+  at: Date;
+  cacheRead: number;
+  input: number;
+  modelCode: number;
+  output: number;
+}
+
+const unknownModel = "gemini-unknown";
+
+const usageField = 9;
+const createdField = 1;
+const generationField = 1;
+const generationUsageField = 4;
+const generationModelField = 19;
+const modelCodeField = 1;
+const inputField = 2;
+const outputField = 3;
+const cacheReadField = 5;
+
+export function antigravityConversationsDir(home: string): string {
+  return resolve(home, ".gemini", "antigravity-cli", "conversations");
+}
+
+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) {
+        return result;
+      }
+      shift += 7n;
+    }
+  };
+  const bytesOf = (length: number): Uint8Array => {
+    if (offset + length > bytes.length) {
+      truncated();
+    }
+    const slice = bytes.subarray(offset, offset + length);
+    offset += length;
+    return slice;
+  };
+  while (offset < bytes.length) {
+    const tag = varint();
+    const number = Number(tag >> 3n);
+    const wireType = Number(tag & 7n);
+    if (wireType === 0) {
+      fields.push({ number, value: varint() });
+    } else if (wireType === 1) {
+      fields.push({ number, value: bytesOf(8) });
+    } else if (wireType === 2) {
+      fields.push({ number, value: bytesOf(Number(varint())) });
+    } else if (wireType === 5) {
+      fields.push({ number, value: bytesOf(4) });
+    } else {
+      throw new Error(`unsupported protobuf wire type ${wireType}`);
+    }
+  }
+  return fields;
+}
+
+function nested(fields: ProtoField[], number: number): ProtoField[] | null {
+  const field = fields.find(
+    (candidate) =>
+      candidate.number === number && candidate.value instanceof Uint8Array,
+  );
+  return field === undefined ? null : decodeMessage(field.value as Uint8Array);
+}
+
+function integer(fields: ProtoField[], number: number): number | null {
+  const field = fields.find(
+    (candidate) =>
+      candidate.number === number && typeof candidate.value === "bigint",
+  );
+  return field === undefined ? null : Number(field.value);
+}
+
+function text(fields: ProtoField[], number: number): string | null {
+  const field = fields.find(
+    (candidate) =>
+      candidate.number === number && candidate.value instanceof Uint8Array,
+  );
+  return field === undefined
+    ? null
+    : new TextDecoder().decode(field.value as Uint8Array);
+}
+
+function openReadOnly(file: string): DatabaseSync {
+  const location = existsSync(`${file}-wal`)
+    ? file
+    : `file:${file.split("/").map(encodeURIComponent).join("/")}?immutable=1`;
+  return new DatabaseSync(location, { readOnly: true });
+}
+
+function readConversation(
+  file: string,
+  names: Map<number, string>,
+  steps: RawStep[],
+): void {
+  const db = openReadOnly(file);
+  try {
+    const generations = db.prepare("SELECT data FROM gen_metadata").all() as {
+      data: Uint8Array;
+    }[];
+    for (const row of generations) {
+      const generation = nested(decodeMessage(row.data), generationField);
+      const usage =
+        generation === null ? null : nested(generation, generationUsageField);
+      const code = usage === null ? null : integer(usage, modelCodeField);
+      const name =
+        generation === null ? null : text(generation, generationModelField);
+      if (code !== null && name !== null) {
+        names.set(code, name);
+      }
+    }
+    const rows = db
+      .prepare("SELECT metadata FROM steps WHERE metadata IS NOT NULL")
+      .all() as { metadata: Uint8Array }[];
+    for (const row of rows) {
+      const fields = decodeMessage(row.metadata);
+      const usage = nested(fields, usageField);
+      const created = nested(fields, createdField);
+      const seconds = created === null ? null : integer(created, 1);
+      if (usage === null || seconds === null) {
+        continue;
+      }
+      steps.push({
+        at: new Date(seconds * 1000),
+        cacheRead: integer(usage, cacheReadField) ?? 0,
+        input: integer(usage, inputField) ?? 0,
+        modelCode: integer(usage, modelCodeField) ?? 0,
+        output: integer(usage, outputField) ?? 0,
+      });
+    }
+  } finally {
+    db.close();
+  }
+}
+
+export async function readAntigravitySteps(
+  dir: string,
+): Promise<AntigravityUsage> {
+  let files: string[];
+  try {
+    files = await readdir(dir);
+  } catch (error) {
+    if ((error as NodeJS.ErrnoException).code === "ENOENT") {
+      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()) {
+    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 {
+    failures,
+    steps: steps.map(({ modelCode, ...step }) => ({
+      ...step,
+      model: names.get(modelCode) ?? unknownModel,
+    })),
+  };
+}

package/src/ccusage.ts

@@ -37,16 +37,25 @@ export function ccusageCliPath(): string {
   return createRequire(import.meta.url).resolve("ccusage/src/cli.js");
 }
 
-export function sinceArgument(today: Date, timezone: string): string {
-  const calendarDate = new Intl.DateTimeFormat("en-CA", {
+export function calendarDate(at: Date, timezone: string): string {
+  return new Intl.DateTimeFormat("en-CA", {
     timeZone: timezone,
     year: "numeric",
     month: "2-digit",
     day: "2-digit",
-  }).format(today);
-  const [year, month, day] = calendarDate.split("-").map(Number);
+  }).format(at);
+}
+
+export function windowStart(today: Date, timezone: string): string {
+  const [year, month, day] = calendarDate(today, timezone)
+    .split("-")
+    .map(Number);
   const start = new Date(Date.UTC(year, month - 1, day - (windowDays - 1)));
-  return start.toISOString().slice(0, 10).replaceAll("-", "");
+  return start.toISOString().slice(0, 10);
+}
+
+export function sinceArgument(today: Date, timezone: string): string {
+  return windowStart(today, timezone).replaceAll("-", "");
 }
 
 export function ccusageArguments(since: string, timezone: string): string[] {

package/src/cli.ts

@@ -1,7 +1,8 @@
 #!/usr/bin/env bun
 import { Command } from "commander";
-import { collect, type Fetcher } from "./collect";
+import { collect } from "./collect";
 import type { CommandRunner } from "./command";
+import type { Fetcher } from "./http";
 import { install } from "./install";
 import { type MachineIdentity, readMachineIdentity } from "./machine";
 import { type CollectorEnv, processEnv } from "./paths";
@@ -35,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

@@ -1,18 +1,22 @@
-import { readCcusageDaily, sinceArgument } from "./ccusage";
+import {
+  antigravityConversationsDir,
+  readAntigravitySteps,
+} from "./antigravity";
+import {
+  calendarDate,
+  readCcusageDaily,
+  sinceArgument,
+  windowStart,
+} from "./ccusage";
 import { type CommandRunner, runCommand } from "./command";
 import { readConfig, runtimeTimezone } from "./config";
+import type { Fetcher } from "./http";
 import { type MachineIdentity, machineId } from "./machine";
-import { mapCcusageDays } from "./mapping";
+import { mapAntigravitySteps, mapCcusageDays } from "./mapping";
 import { type CollectorEnv, collectorPaths, processEnv } from "./paths";
+import { loadPrices } from "./pricing";
 import type { UsageDay, UsageReport } from "./usage";
 
-export interface HttpResponse {
-  status: number;
-  text(): Promise<string>;
-}
-
-export type Fetcher = (url: string, init: RequestInit) => Promise<HttpResponse>;
-
 export interface CollectOptions {
   identity: MachineIdentity;
   env?: CollectorEnv;
@@ -23,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 };
 
@@ -52,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),
@@ -80,11 +90,34 @@ 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(
+  home: string,
+  today: Date,
+  timezone: string,
+  fetcher: Fetcher,
+  pricesFile: string,
+): Promise<AntigravityRows> {
+  const since = windowStart(today, timezone);
+  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 { days: [], warnings };
+  }
+  const prices = await loadPrices(fetcher, pricesFile, today);
+  return { days: mapAntigravitySteps(steps, timezone, prices), warnings };
 }
 
 export async function collect(options: CollectOptions): Promise<CollectResult> {
-  const paths = collectorPaths(options.env ?? processEnv());
+  const env = options.env ?? processEnv();
+  const paths = collectorPaths(env);
   const config = await readConfig(paths.configFile);
   if (config.kind === "missing") {
     return { configFile: paths.configFile, kind: "missing-config" };
@@ -100,28 +133,45 @@ export async function collect(options: CollectOptions): Promise<CollectResult> {
   const timezone =
     config.config.timezone ?? options.timezone ?? runtimeTimezone();
   const runner = options.runner ?? runCommand;
+  const fetcher = options.fetcher ?? fetch;
+  const today = options.today ?? new Date();
 
   let days: UsageDay[];
   try {
     const daily = await readCcusageDaily(
       runner,
-      sinceArgument(options.today ?? new Date(), timezone),
+      sinceArgument(today, timezone),
       timezone,
     );
     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(
-      options.fetcher ?? fetch,
+      fetcher,
       config.config.url,
       config.config.key,
       { days, machine, timezone },
+      antigravity.warnings,
     );
   } catch (error) {
     return { kind: "failed", message: messageOf(error) };

package/src/http.ts

@@ -0,0 +1,6 @@
+export interface HttpResponse {
+  status: number;
+  text(): Promise<string>;
+}
+
+export type Fetcher = (url: string, init: RequestInit) => Promise<HttpResponse>;

package/src/mapping.ts

@@ -1,6 +1,10 @@
-import type { CcusageDaily } from "./ccusage";
+import type { AntigravityStep } from "./antigravity";
+import { type CcusageDaily, calendarDate } from "./ccusage";
+import { costOf, type PriceTable } from "./pricing";
 import type { UsageDay } from "./usage";
 
+export const antigravityProvider = "antigravity";
+
 const agentModelPrefix = /^\[[^\]]*\]\s*/;
 
 const rowKey = (day: UsageDay): string =>
@@ -46,6 +50,57 @@ export function mapCcusageDays(output: CcusageDaily): UsageDay[] {
     }
   }
 
+  return sortedRows(rows);
+}
+
+export function mapAntigravitySteps(
+  steps: AntigravityStep[],
+  timezone: string,
+  prices: PriceTable,
+): UsageDay[] {
+  const rows = new Map<string, UsageDay>();
+
+  for (const step of steps) {
+    const row: UsageDay = {
+      cache_create: 0,
+      cache_read: step.cacheRead,
+      cost_usd: 0,
+      date: calendarDate(step.at, timezone),
+      input: step.input,
+      model: step.model,
+      output: step.output,
+      provider: antigravityProvider,
+    };
+    if (isEmptyRow(row)) {
+      continue;
+    }
+    const key = rowKey(row);
+    const existing = rows.get(key);
+    if (existing === undefined) {
+      rows.set(key, row);
+      continue;
+    }
+    existing.cache_read += row.cache_read;
+    existing.input += row.input;
+    existing.output += row.output;
+  }
+
+  for (const row of rows.values()) {
+    const price = prices.get(row.model) ?? prices.get(`gemini/${row.model}`);
+    row.cost_usd =
+      price === undefined
+        ? 0
+        : costOf(price, {
+            cacheRead: row.cache_read,
+            input: row.input,
+            output: row.output,
+          });
+  }
+
+  return sortedRows(rows);
+}
+
+function sortedRows(rows: Map<string, UsageDay>): UsageDay[] {
   return [...rows.values()].sort(
     (a, b) =>
       a.date.localeCompare(b.date) ||

package/src/paths.ts

@@ -12,6 +12,7 @@ export interface CollectorEnv {
 export interface CollectorPaths {
   configFile: string;
   plist: string;
+  pricesFile: string;
   service: string;
   stderrLog: string;
   stdoutLog: string;
@@ -54,6 +55,7 @@ export function collectorPaths(env: CollectorEnv): CollectorPaths {
       "LaunchAgents",
       "dev.tokenmax.collector.plist",
     ),
+    pricesFile: resolve(configHome, "tokenmax", "litellm-prices.json"),
     service: resolve(home, ".config", "systemd", "user", "tokenmax.service"),
     stderrLog: resolve(logDir, "tokenmax.err.log"),
     stdoutLog: resolve(logDir, "tokenmax.log"),

package/src/pricing.ts

@@ -0,0 +1,120 @@
+import { mkdir, readFile, rename, stat, writeFile } from "node:fs/promises";
+import { dirname } from "node:path";
+import { z } from "zod";
+import type { Fetcher } from "./http";
+
+export interface ModelPrice {
+  cacheRead: number;
+  input: number;
+  output: number;
+}
+
+export type PriceTable = Map<string, ModelPrice>;
+
+export interface PricedUsage {
+  cacheRead: number;
+  input: number;
+  output: number;
+}
+
+export const litellmPricesUrl =
+  "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json";
+
+const maxAgeMs = 24 * 60 * 60 * 1000;
+
+const pricedEntry = z.object({
+  cache_read_input_token_cost: z.number().optional(),
+  input_cost_per_token: z.number(),
+  output_cost_per_token: z.number(),
+});
+
+const litellmPrices = z.record(z.string(), z.unknown());
+
+export function parseLitellmPrices(source: string): PriceTable {
+  const table: PriceTable = new Map();
+  for (const [model, entry] of Object.entries(
+    litellmPrices.parse(JSON.parse(source)),
+  )) {
+    const priced = pricedEntry.safeParse(entry);
+    if (priced.success) {
+      table.set(model, {
+        cacheRead: priced.data.cache_read_input_token_cost ?? 0,
+        input: priced.data.input_cost_per_token,
+        output: priced.data.output_cost_per_token,
+      });
+    }
+  }
+  return table;
+}
+
+export function costOf(price: ModelPrice, usage: PricedUsage): number {
+  return (
+    usage.input * price.input +
+    usage.output * price.output +
+    usage.cacheRead * price.cacheRead
+  );
+}
+
+async function readCache(
+  pricesFile: string,
+): Promise<{ modifiedAt: number; table: PriceTable } | null> {
+  let source: string;
+  let modifiedAt: number;
+  try {
+    [source, modifiedAt] = await Promise.all([
+      readFile(pricesFile, "utf8"),
+      stat(pricesFile).then((info) => info.mtimeMs),
+    ]);
+  } 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> {
+  const response = await fetcher(litellmPricesUrl, { method: "GET" });
+  const source = await response.text();
+  if (response.status !== 200) {
+    throw new Error(`LiteLLM responded ${response.status}`);
+  }
+  return source;
+}
+
+export async function loadPrices(
+  fetcher: Fetcher,
+  pricesFile: string,
+  now: Date,
+): Promise<PriceTable> {
+  const cached = await readCache(pricesFile);
+  if (cached !== null && now.getTime() - cached.modifiedAt < maxAgeMs) {
+    return cached.table;
+  }
+  let source: string;
+  let table: PriceTable;
+  try {
+    source = await fetchPrices(fetcher);
+    table = parseLitellmPrices(source);
+  } catch (error) {
+    if (cached !== null) {
+      return cached.table;
+    }
+    const message = error instanceof Error ? error.message : String(error);
+    throw new Error(`could not load model prices: ${message}`);
+  }
+  await writeCache(pricesFile, source);
+  return table;
+}