tokentracker-cli 1.1.3 → 1.1.4

package/src/lib/trae-sqlite.js

@@ -0,0 +1,143 @@
+const fs = require("node:fs");
+
+let enginePromise;
+let queryQueue = Promise.resolve();
+
+async function createEngine() {
+  const [{ default: initialize }, sqlite, { FacadeVFS }] = await Promise.all([
+    import("@journeyapps/wa-sqlite/dist/wa-sqlite.mjs"),
+    import("@journeyapps/wa-sqlite"),
+    import("@journeyapps/wa-sqlite/src/FacadeVFS.js"),
+  ]);
+  const module = await initialize({
+    wasmBinary: fs.readFileSync(require.resolve("@journeyapps/wa-sqlite/dist/wa-sqlite.wasm")),
+  });
+  const api = sqlite.Factory(module);
+  const snapshots = new Map();
+  const handles = new Map();
+
+  // Register one VFS for the lifetime of the engine. Each query gets a distinct
+  // virtual filename, without accumulating WASM callbacks or touching disk.
+  class TraeVFS extends FacadeVFS {
+    constructor() {
+      super("tokentracker-trae", module);
+    }
+
+    jOpen(name, fileId, flags, outputFlags) {
+      const state = snapshots.get(name);
+      if (!state || !(flags & sqlite.SQLITE_OPEN_MAIN_DB)) return sqlite.SQLITE_CANTOPEN;
+      handles.set(fileId, state);
+      outputFlags.setInt32(0, sqlite.SQLITE_OPEN_READONLY, true);
+      return sqlite.SQLITE_OK;
+    }
+
+    jRead(fileId, output, offset) {
+      const state = handles.get(fileId);
+      if (!state) return sqlite.SQLITE_IOERR_READ;
+      // FacadeVFS passes a proxy over WASM memory, not a native Uint8Array.
+      // Keep the temporary buffer small and erase it after copying its page.
+      const buffer = Buffer.alloc(output.byteLength);
+      try {
+        const count = state.snapshot.read(buffer, offset);
+        if (!Number.isSafeInteger(count) || count < 0 || count > buffer.length) {
+          throw new Error("Invalid TRAE snapshot read length.");
+        }
+        buffer.fill(0, count);
+        output.set(buffer);
+        return count === buffer.length ? sqlite.SQLITE_OK : sqlite.SQLITE_IOERR_SHORT_READ;
+      } catch (err) {
+        state.error ||= err;
+        output.fill(0);
+        return sqlite.SQLITE_IOERR_READ;
+      } finally {
+        buffer.fill(0);
+      }
+    }
+
+    jFileSize(fileId, output) {
+      const state = handles.get(fileId);
+      if (!state) return sqlite.SQLITE_IOERR_FSTAT;
+      output.setBigInt64(0, BigInt(state.snapshot.byteLength), true);
+      return sqlite.SQLITE_OK;
+    }
+
+    jClose(fileId) {
+      handles.delete(fileId);
+      return sqlite.SQLITE_OK;
+    }
+
+    jAccess(name, flags, output) {
+      output.setInt32(0, snapshots.has(name) ? 1 : 0, true);
+      return sqlite.SQLITE_OK;
+    }
+
+    jWrite() { return sqlite.SQLITE_READONLY; }
+    jTruncate() { return sqlite.SQLITE_READONLY; }
+    jDelete() { return sqlite.SQLITE_READONLY; }
+  }
+
+  const vfs = new TraeVFS();
+  api.vfs_register(vfs, false);
+  return { api, sqlite, snapshots, handles, vfs, nextId: 0 };
+}
+
+/**
+ * Query a read-only, random-access snapshot without materializing its database.
+ * The caller owns snapshot cleanup; only the projected SQL rows leave this VFS.
+ */
+async function querySnapshot(snapshot, callback) {
+  if (!Number.isSafeInteger(snapshot.byteLength) || snapshot.byteLength < 0) {
+    throw new Error("Invalid TRAE snapshot size.");
+  }
+  enginePromise ||= createEngine().catch((err) => {
+    enginePromise = undefined;
+    throw err;
+  });
+  const engine = await enginePromise;
+  const { api, sqlite, snapshots, handles, vfs } = engine;
+  const name = `trae-${++engine.nextId}.db`;
+  const state = { snapshot, error: null };
+  snapshots.set(name, state);
+  let db;
+  let failure;
+  try {
+    db = await api.open_v2(name, sqlite.SQLITE_OPEN_READONLY, vfs.name);
+    const query = async (sql) => {
+      const rows = [];
+      try {
+        await api.exec(db, sql, (values, columns) => {
+          rows.push(Object.fromEntries(columns.map((column, index) => [column, values[index]])));
+        });
+      } catch (err) {
+        throw state.error || err;
+      }
+      return rows;
+    };
+    await query("PRAGMA query_only=ON; PRAGMA cache_size=-4096; PRAGMA temp_store=MEMORY;");
+    return await callback(query);
+  } catch (err) {
+    failure = state.error || err;
+    throw failure;
+  } finally {
+    try {
+      if (db) await api.close(db);
+    } catch (err) {
+      if (!failure) throw state.error || err;
+    } finally {
+      snapshots.delete(name);
+      for (const [fileId, opened] of handles) {
+        if (opened === state) handles.delete(fileId);
+      }
+    }
+  }
+}
+
+function withTraeSqlite(snapshot, callback) {
+  // The package's async convenience API shares temporary WASM allocations.
+  // Serialize access to the single engine even when callers read in parallel.
+  const operation = queryQueue.then(() => querySnapshot(snapshot, callback));
+  queryQueue = operation.catch(() => {});
+  return operation;
+}
+
+module.exports = { withTraeSqlite };

package/src/lib/trae-usage.js

@@ -0,0 +1,121 @@
+"use strict";
+
+// TRAE stores the last request's usage alongside optional whole-turn prompt
+// and completion counters. Reconcile the request first: its total cannot
+// validate a whole-turn count, and its cache count is not a whole-turn count.
+function normalizeTraeUsage(raw, { model } = {}) {
+  let usage = raw;
+  if (typeof usage === "string") {
+    try { usage = JSON.parse(usage); } catch { return null; }
+  }
+  if (!usage || typeof usage !== "object" || Array.isArray(usage)) return null;
+  const fields = [
+    "prompt_tokens", "prompt_tokens_total", "input_tokens",
+    "completion_tokens", "completion_tokens_total", "output_tokens",
+    "cache_read_input_tokens", "cache_creation_input_tokens",
+    "reasoning_tokens", "reasoning_output_tokens", "total_tokens",
+  ];
+  for (const field of fields) {
+    if (usage[field] != null && (!Number.isSafeInteger(usage[field]) || usage[field] < 0)) {
+      return null;
+    }
+  }
+  for (const [left, right] of [["prompt_tokens", "input_tokens"],
+    ["completion_tokens", "output_tokens"], ["reasoning_tokens", "reasoning_output_tokens"]]) {
+    if (usage[left] != null && usage[right] != null && usage[left] !== usage[right]) return null;
+  }
+  if (!["prompt_tokens", "prompt_tokens_total", "input_tokens", "completion_tokens",
+    "completion_tokens_total", "output_tokens"].some((field) => usage[field] != null)) {
+    return null;
+  }
+  const input = usage.prompt_tokens ?? usage.input_tokens ?? usage.prompt_tokens_total ?? 0;
+  const output = usage.completion_tokens ?? usage.output_tokens ?? usage.completion_tokens_total ?? 0;
+  const aggregateInput = usage.prompt_tokens_total > 0 ? usage.prompt_tokens_total : input;
+  const aggregateOutput = usage.completion_tokens_total > 0 ? usage.completion_tokens_total : output;
+  if (aggregateInput < input || aggregateOutput < output) return null;
+  const cached = usage.cache_read_input_tokens ?? 0;
+  let written = usage.cache_creation_input_tokens ?? 0;
+  let reasoning = usage.reasoning_tokens ?? usage.reasoning_output_tokens ?? 0;
+  let estimated = false;
+  const isGemini = /^gemini(?:-|$)/i.test(normalizeTraeModel(model ?? usage.name));
+  // Some TRAE Gemini rows copy the cache-read counter into the write field,
+  // even when their sum exceeds the entire prompt. Gemini's GenerateContent
+  // metadata reports cachedContentTokenCount, not an additional write count.
+  // Keep the reported read and mark this narrow compatibility repair estimated.
+  if (isGemini && cached > 0 && written === cached && cached + written > input) {
+    written = 0;
+    estimated = true;
+  }
+  // Older TRAE Gemini rows omit thoughts while retaining Gemini's total
+  // (prompt + candidates + thoughts). Infer only that documented residual,
+  // never an arbitrary provider's unexplained total.
+  if (isGemini && usage.reasoning_tokens == null && usage.reasoning_output_tokens == null &&
+    usage.total_tokens > input + output && written === 0 && cached <= input) {
+    reasoning = usage.total_tokens - input - output;
+    estimated = true;
+  }
+  // A total is necessary to distinguish inclusive/exclusive counters when
+  // cache or reasoning is present. Missing metadata must not imply $0 usage
+  // or guessed input. A zero *_total field means unset, not an empty turn.
+  if (usage.total_tokens == null && (cached || written || reasoning)) return null;
+  const total = usage.total_tokens ?? (input + output);
+  if (!Number.isSafeInteger(total)) return null;
+  const candidates = new Map();
+  for (const includedCache of new Set([cached + written, cached, 0])) {
+    for (const includedReasoning of new Set([reasoning, 0])) {
+      const noncached = input - includedCache;
+      const nonreasoning = output - includedReasoning;
+      if (noncached < 0 || nonreasoning < 0) continue;
+      if (noncached + cached + written + nonreasoning + reasoning !== total) continue;
+      const totals = {
+        input_tokens: noncached,
+        cached_input_tokens: cached,
+        cache_creation_input_tokens: written,
+        output_tokens: nonreasoning,
+        reasoning_output_tokens: reasoning,
+        total_tokens: total,
+      };
+      candidates.set(JSON.stringify(totals), totals);
+    }
+  }
+  if (candidates.size !== 1) return null;
+  const totals = candidates.values().next().value;
+  if (aggregateInput !== input || aggregateOutput !== output) {
+    // The *_total counters cover the whole turn; the cache and reasoning fields
+    // describe its final request. The difference in prompt tokens has no cache
+    // split, so it counts toward total_tokens but no priced column. Completion
+    // tokens are priced the same with or without caching and stay in output.
+    const unpriced = aggregateInput - input;
+    totals.output_tokens += aggregateOutput - output;
+    totals.total_tokens += unpriced + aggregateOutput - output;
+    if (unpriced > 0) totals.unpriced_input_tokens = unpriced;
+  }
+  if (Object.values(totals).some((value) => !Number.isSafeInteger(value))) return null;
+  if (estimated) totals.usage_precision = "estimated";
+  return totals;
+}
+
+function normalizeTraeModel(value) {
+  if (typeof value !== "string") return "trae-unknown";
+  const model = value.trim().split("__")[0].toLowerCase().replace(/\s+/g, "-");
+  // TRAE's display alias otherwise fuzzy-matches a regional Sonnet SKU and
+  // loses cache pricing. Sonnet 4 has the fixed 2025-05-14 API model id.
+  if (model === "claude-4-sonnet") return "claude-sonnet-4-20250514";
+  return model && model !== "-" && !model.includes("|") ? model : "trae-unknown";
+}
+
+function traeTimestamp(value) {
+  if (value == null || value === "") return null;
+  let millis;
+  if (typeof value === "number" || (typeof value === "string" && /^\d+(\.\d+)?$/.test(value))) {
+    const numeric = Number(value);
+    millis = numeric < 100_000_000_000 ? numeric * 1000 : numeric;
+  } else if (typeof value === "string") {
+    millis = Date.parse(value);
+  }
+  if (!Number.isFinite(millis) || millis <= 0) return null;
+  const date = new Date(millis);
+  return Number.isFinite(date.getTime()) ? date.toISOString() : null;
+}
+
+module.exports = { normalizeTraeUsage, normalizeTraeModel, traeTimestamp };