tokenmax-collector 0.1.0 → 0.2.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "tokenmax-collector",
-  "version": "0.1.0",
+  "version": "0.2.0",
   "description": "Report local ccusage token usage to a tokenmax instance",
   "type": "module",
   "bin": {

package/src/antigravity.ts

@@ -0,0 +1,181 @@
+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;
+}
+
+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 varint = (): bigint => {
+    let result = 0n;
+    let shift = 0n;
+    for (;;) {
+      const byte = bytes[offset++];
+      result |= BigInt(byte & 0x7f) << shift;
+      if (byte < 0x80) {
+        return result;
+      }
+      shift += 7n;
+    }
+  };
+  const bytesOf = (length: number): Uint8Array => {
+    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}?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<AntigravityStep[]> {
+  let files: string[];
+  try {
+    files = await readdir(dir);
+  } catch (error) {
+    if ((error as NodeJS.ErrnoException).code === "ENOENT") {
+      return [];
+    }
+    throw error;
+  }
+  const names = new Map<number, string>();
+  const steps: RawStep[] = [];
+  for (const file of files.filter((name) => name.endsWith(".db")).sort()) {
+    readConversation(resolve(dir, file), names, steps);
+  }
+  return 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";

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;
@@ -83,8 +87,27 @@ async function report(
   return { accepted, kind: "reported", machine: usage.machine };
 }
 
+async function antigravityDays(
+  home: string,
+  today: Date,
+  timezone: string,
+  fetcher: Fetcher,
+  pricesFile: string,
+): Promise<UsageDay[]> {
+  const since = windowStart(today, timezone);
+  const steps = (
+    await readAntigravitySteps(antigravityConversationsDir(home))
+  ).filter((step) => calendarDate(step.at, timezone) >= since);
+  if (steps.length === 0) {
+    return [];
+  }
+  const prices = await loadPrices(fetcher, pricesFile, today);
+  return mapAntigravitySteps(steps, timezone, prices);
+}
+
 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,15 +123,26 @@ 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);
+    days = [
+      ...mapCcusageDays(daily),
+      ...(await antigravityDays(
+        env.home,
+        today,
+        timezone,
+        fetcher,
+        paths.pricesFile,
+      )),
+    ];
   } catch (error) {
     return { kind: "failed", message: messageOf(error) };
   }
@@ -117,12 +151,11 @@ export async function collect(options: CollectOptions): Promise<CollectResult> {
   }
 
   try {
-    return await report(
-      options.fetcher ?? fetch,
-      config.config.url,
-      config.config.key,
-      { days, machine, timezone },
-    );
+    return await report(fetcher, config.config.url, config.config.key, {
+      days,
+      machine,
+      timezone,
+    });
   } 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,108 @@
+import { mkdir, readFile, 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; source: string } | null> {
+  try {
+    const [source, info] = await Promise.all([
+      readFile(pricesFile, "utf8"),
+      stat(pricesFile),
+    ]);
+    return { modifiedAt: info.mtimeMs, source };
+  } catch (error) {
+    if ((error as NodeJS.ErrnoException).code === "ENOENT") {
+      return null;
+    }
+    throw error;
+  }
+}
+
+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 parseLitellmPrices(cached.source);
+  }
+  let source: string;
+  let table: PriceTable;
+  try {
+    source = await fetchPrices(fetcher);
+    table = parseLitellmPrices(source);
+  } catch (error) {
+    if (cached !== null) {
+      return parseLitellmPrices(cached.source);
+    }
+    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);
+  return table;
+}