tokentracker-cli 0.93.3 → 0.93.5

package/src/lib/rollout.js

@@ -2698,7 +2698,8 @@ async function parseOpencodeMessageFile({
     };
   }
 
-  const model = normalizeModelInput(msg?.modelID || msg?.model || msg?.modelId) || DEFAULT_MODEL;
+  const { modelId: fileModelId } = normalizeOpencodeModelFields(msg);
+  const model = fileModelId || DEFAULT_MODEL;
   const bucket = getHourlyBucket(hourlyState, source, model, bucketStart);
   addTotals(bucket.totals, delta);
   touchedBuckets.add(bucketKey(source, model, bucketStart));
@@ -3289,8 +3290,11 @@ function deriveOpencodeMessageFingerprint({ msg, totals, source }) {
   const created = coerceEpochMs(msg?.time?.created) || 0;
   const completed = coerceEpochMs(msg?.time?.completed) || 0;
   if (!created && !completed) return null;
-  const model = normalizeModelInput(msg?.modelID || msg?.model || msg?.modelId) || DEFAULT_MODEL;
-  const provider = normalizeMessageKeyPart(msg?.providerID || msg?.provider || msg?.providerId);
+  // v1 keeps flat modelID/providerID strings; v2 nests them in
+  // `model: { id, providerID }` — normalize both (see normalizeOpencodeModelFields).
+  const { modelId, providerId } = normalizeOpencodeModelFields(msg);
+  const model = modelId || DEFAULT_MODEL;
+  const provider = providerId;
   const raw = [
     normalizeSourceInput(source) || "opencode",
     created,
@@ -3402,7 +3406,8 @@ function repairCountedOpencodeForkCopy({
   if (!timestampMs) return false;
   const bucketStart = toUtcHalfHourStart(new Date(timestampMs).toISOString());
   if (!bucketStart) return false;
-  const model = normalizeModelInput(msg?.modelID || msg?.model || msg?.modelId) || DEFAULT_MODEL;
+  const { modelId: repairModelId } = normalizeOpencodeModelFields(msg);
+  const model = repairModelId || DEFAULT_MODEL;
   const counted = { ...totals, conversation_count: 1 };
   const bucket = getHourlyBucket(hourlyState, source, model, bucketStart);
   subtractTotals(bucket.totals, counted);
@@ -3635,6 +3640,33 @@ function normalizeModelInput(value) {
   return trimmed.length > 0 ? trimmed : null;
 }
 
+// OpenCode v1 stores the model as flat strings on every assistant message
+// (`modelID`, `providerID`), while OpenCode v2 (the `opencode2` beta, which
+// moved messages into the `session_message` table) nests them inside a single
+// object: `model: { id, providerID, variant }`. Some forks also wrote a plain
+// string into `model`. Resolve both shapes to { modelId, providerId } so
+// bucket keys and fork fingerprints stay identical across versions.
+function normalizeOpencodeModelFields(msg) {
+  const directModel =
+    normalizeModelInput(msg?.modelID) ||
+    normalizeModelInput(typeof msg?.model === "string" ? msg.model : null) ||
+    normalizeModelInput(msg?.modelId);
+  if (directModel) {
+    return {
+      modelId: directModel,
+      providerId: normalizeMessageKeyPart(msg?.providerID || msg?.provider || msg?.providerId),
+    };
+  }
+  const nested = msg?.model;
+  if (nested && typeof nested === "object") {
+    return {
+      modelId: normalizeModelInput(nested.id),
+      providerId: normalizeMessageKeyPart(nested.providerID),
+    };
+  }
+  return { modelId: null, providerId: "" };
+}
+
 async function resolveProjectMetaForPath(startDir, cache) {
   if (!startDir || typeof startDir !== "string") return null;
   if (cache && cache.has(startDir)) return cache.get(startDir);
@@ -4298,12 +4330,79 @@ async function walkOpencodeMessages(dir, out) {
 }
 
 // ---------------------------------------------------------------------------
-// OpenCode SQLite DB reader (v1.2+ stores messages in opencode.db)
+// OpenCode SQLite DB reader
+//
+// Four real database shapes exist in the wild (see PR #501 / opencode2):
+//   A  pure v1:        `message`(data) + `session_message`(empty) [+ `session`]
+//   B  beta-17887:     `session_message`(data) + `session_v2`
+//   C  upstream v2:    `session_message`(data) + `session`
+//   D  transitional:   `message`(old data) + `session_message`(new data)
+//
+// The reader must serve all four without firing a doomed SQL at type A (an
+// empty `session_message` table still EXISTS — querying it with a JOIN against
+// `session_v2` that does not exist throws "no such table" every sync, and the
+// old code silently swallowed that error). The combined probe below answers
+// both "is there v2 data?" and "which session table exists?" in one round-trip.
 // ---------------------------------------------------------------------------
 
+const OPENCODE_DB_V1_MESSAGE_SQL =
+  `SELECT id, session_id, time_updated, data FROM message ` +
+  `WHERE json_extract(data, '$.role') = 'assistant' ORDER BY time_created ASC`;
+
+// Build the v2 query. When a session table exists the LEFT JOIN restores the
+// project directory for downstream attribution; when it does not (type B
+// without session_v2, or an exotic fork) the join and directory column are
+// omitted and tokens are counted as-is.
+function buildV2Sql(sessionTable) {
+  const joinClause = sessionTable
+    ? `LEFT JOIN ${sessionTable} s ON s.id = sm.session_id `
+    : "";
+  const directorySelect = sessionTable
+    ? `s.directory AS directory, `
+    : "";
+  return (
+    `SELECT sm.id AS id, sm.session_id AS session_id, sm.time_updated AS time_updated, ` +
+    `${directorySelect}sm.data AS data ` +
+    `FROM session_message sm ${joinClause}` +
+    `WHERE sm.type = 'assistant' ORDER BY sm.time_created ASC`
+  );
+}
+
+// One combined probe: hasRows (is there any data in session_message?) plus
+// sessionTable (which session table exists, session_v2 preferred over session
+// via DESC). Returns null when the probe itself fails — callers treat that as
+// v1-only, matching the pre-existing silent-degradation contract.
+function detectOpencodeMessageLayout(dbPath, sqliteOptions = {}) {
+  let rows;
+  try {
+    rows = readSqliteJsonRows(
+      dbPath,
+      `SELECT (SELECT 1 FROM session_message LIMIT 1) AS hasRows,
+              (SELECT name FROM sqlite_master WHERE name IN ('session','session_v2') ORDER BY name DESC LIMIT 1) AS sessionTable`,
+      { label: "OpenCode", timeout: 10_000, maxBuffer: 1024 * 1024, ...sqliteOptions },
+    );
+  } catch (_e) {
+    return null;
+  }
+  if (!Array.isArray(rows) || rows.length === 0) return null;
+  const row = rows[0];
+  return {
+    hasRows: Boolean(row?.hasRows),
+    sessionTable: typeof row?.sessionTable === "string" && row.sessionTable.trim()
+      ? row.sessionTable.trim()
+      : null,
+  };
+}
+
+// Shared provider resolver: v1 rows carry a top-level providerID, v2 rows nest
+// it under model.providerID. Used by the mimo/zcode discriminators so they work
+// against both generations without duplicating the resolution logic.
+function opencodeMessageProvider(data) {
+  return data?.providerID || data?.model?.providerID || "";
+}
+
 function readOpencodeDbMessages(dbPath, sqliteOptions = {}) {
   if (!dbPath || !fssync.existsSync(dbPath)) return [];
-  const sql = `SELECT id, session_id, time_updated, data FROM message WHERE json_extract(data, '$.role') = 'assistant' ORDER BY time_created ASC`;
 
   let snapshot = null;
   let effectiveDbPath = dbPath;
@@ -4314,14 +4413,16 @@ function readOpencodeDbMessages(dbPath, sqliteOptions = {}) {
     } catch (_e) { }
   }
 
-  try {
-    const rows = readSqliteJsonRows(effectiveDbPath, sql, {
-      label: "OpenCode",
-      maxBuffer: 50 * 1024 * 1024,
-      timeout: 30_000,
-      ...sqliteOptions,
-    });
-    const out = [];
+  // Parse raw SQL rows into the shared { id, sessionID, timeUpdated, data }
+  // shape; optionally restore v2's session-level project directory as the
+  // per-message `path.cwd` the rest of the pipeline expects.
+  //
+  // D-state (transitional) union: both the legacy `message` table and the new
+  // `session_message` table may carry rows. Cross-table duplicates are caught
+  // by the existing cursor-key + #426 fingerprint dedup in
+  // parseOpencodeDbIncremental, so unioning here is safe — the downstream
+  // parser collapses identical sessionID|messageID pairs before bucketing.
+  const appendRows = (rows, isV2, out) => {
     for (const row of rows) {
       if (!row || typeof row.data !== "string") continue;
       let data;
@@ -4338,6 +4439,9 @@ function readOpencodeDbMessages(dbPath, sqliteOptions = {}) {
         toNonNegativeInt(tokens.output) > 0 ||
         toNonNegativeInt(tokens.reasoning) > 0;
       if (!hasTokens) continue;
+      if (isV2 && typeof row.directory === "string" && row.directory.trim()) {
+        data.path = { ...(typeof data.path === "object" && data.path ? data.path : {}), cwd: row.directory };
+      }
       out.push({
         id: row.id || data.id,
         sessionID: row.session_id || data.sessionID,
@@ -4345,6 +4449,39 @@ function readOpencodeDbMessages(dbPath, sqliteOptions = {}) {
         data,
       });
     }
+  };
+
+  const readGeneration = (sql) => {
+    try {
+      const rows = readSqliteJsonRows(effectiveDbPath, sql, {
+        label: "OpenCode",
+        maxBuffer: 50 * 1024 * 1024,
+        timeout: 30_000,
+        ...sqliteOptions,
+        throwOnReadFailure: true,
+      });
+      return Array.isArray(rows) ? rows : [];
+    } catch (_e) {
+      return null; // query-level failure (e.g. wrong generation's table)
+    }
+  };
+
+  try {
+    // Combined probe drives the orchestration: v1 always runs (it is the
+    // baseline every database carries or degrades to), v2 runs only when the
+    // probe confirms session_message has rows — this avoids firing a JOIN
+    // against a non-existent session table on type-A databases.
+    const probe = detectOpencodeMessageLayout(effectiveDbPath, sqliteOptions);
+    const out = [];
+
+    const rows1 = readGeneration(OPENCODE_DB_V1_MESSAGE_SQL);
+    if (rows1) appendRows(rows1, false, out);
+
+    if (probe?.hasRows) {
+      const rows2 = readGeneration(buildV2Sql(probe.sessionTable));
+      if (rows2) appendRows(rows2, true, out);
+    }
+
     return out;
   } finally {
     if (snapshot) snapshot.cleanup();
@@ -4369,7 +4506,7 @@ function readOpencodeDbMessages(dbPath, sqliteOptions = {}) {
 // subsumes it.
 function isMimoNativeMessage(data) {
   if (!data) return false;
-  const provider = String(data.providerID || "").toLowerCase();
+  const provider = String(opencodeMessageProvider(data)).toLowerCase();
   return provider === "mimo" || provider === "xiaomi";
 }
 
@@ -4403,7 +4540,7 @@ function readMimoDbMessages(dbPath, sqliteOptions = {}) {
 // re-count it. Mirrors the Mimo discipline.
 function isZcodeNativeMessage(data) {
   if (!data) return false;
-  const provider = String(data.providerID || "").toLowerCase();
+  const provider = String(opencodeMessageProvider(data)).toLowerCase();
   if (!provider) return false;
   return !(
     provider.includes("anthropic") ||
@@ -4412,10 +4549,195 @@ function isZcodeNativeMessage(data) {
   );
 }
 
+const ZCODE_NATIVE_USAGE_COLUMNS = new Set([
+  "id",
+  "logical_request_id",
+  "attempt_index",
+  "session_id",
+  "provider_id",
+  "model_id",
+  "status",
+  "started_at",
+  "input_tokens",
+  "output_tokens",
+  "reasoning_tokens",
+  "cache_creation_input_tokens",
+  "cache_read_input_tokens",
+]);
+
+function detectZcodeNativeUsageLayout(dbPath, sqliteOptions = {}) {
+  if (!dbPath || !fssync.existsSync(dbPath)) return null;
+  let rows;
+  try {
+    rows = readSqliteJsonRows(
+      dbPath,
+      `SELECT 'model_usage' AS table_name, name FROM pragma_table_info('model_usage')
+       UNION ALL
+       SELECT 'session' AS table_name, name FROM pragma_table_info('session')`,
+      {
+        label: "ZCode",
+        timeout: 10_000,
+        maxBuffer: 1024 * 1024,
+        ...sqliteOptions,
+        throwOnReadFailure: true,
+      },
+    );
+  } catch (_error) {
+    return null;
+  }
+  const modelUsageColumns = new Set(
+    rows
+      .filter((row) => !row?.table_name || row.table_name === "model_usage")
+      .map((row) => String(row?.name || "")),
+  );
+  if (![...ZCODE_NATIVE_USAGE_COLUMNS].every((name) => modelUsageColumns.has(name))) {
+    return null;
+  }
+  const sessionColumns = new Set(
+    rows
+      .filter((row) => row?.table_name === "session")
+      .map((row) => String(row?.name || "")),
+  );
+  return {
+    hasSessionDirectory: sessionColumns.has("id") && sessionColumns.has("directory"),
+  };
+}
+
+function hasZcodeNativeUsageSchema(dbPath, sqliteOptions = {}) {
+  if (!dbPath || !fssync.existsSync(dbPath)) return false;
+  let snapshot = null;
+  let effectiveDbPath = dbPath;
+  if (isUncPath(dbPath)) {
+    try {
+      snapshot = snapshotSqliteDb(dbPath);
+      effectiveDbPath = snapshot.path;
+    } catch (_error) {
+      // Fall through to the direct read: some UNC servers support SQLite's
+      // read-only locking semantics even when the WSL bridge does not.
+    }
+  }
+  try {
+    return detectZcodeNativeUsageLayout(effectiveDbPath, sqliteOptions) !== null;
+  } finally {
+    if (snapshot) snapshot.cleanup();
+  }
+}
+
+function buildZcodeNativeUsageSql({ hasSessionDirectory }) {
+  const directorySelect = hasSessionDirectory ? ", s.directory AS directory" : "";
+  const directoryJoin = hasSessionDirectory
+    ? " LEFT JOIN session AS s ON s.id = mu.session_id"
+    : "";
+  return `SELECT
+    mu.id,
+    mu.logical_request_id,
+    mu.attempt_index,
+    mu.session_id,
+    mu.provider_id,
+    mu.model_id,
+    mu.started_at,
+    mu.input_tokens,
+    mu.output_tokens,
+    mu.reasoning_tokens,
+    mu.cache_creation_input_tokens,
+    mu.cache_read_input_tokens
+    ${directorySelect}
+    FROM model_usage AS mu${directoryJoin}
+    WHERE mu.status = 'completed'
+      AND trim(mu.model_id) != ''
+      AND (
+        mu.input_tokens > 0 OR mu.output_tokens > 0 OR mu.reasoning_tokens > 0 OR
+        mu.cache_creation_input_tokens > 0 OR mu.cache_read_input_tokens > 0
+      )
+    ORDER BY mu.started_at ASC, mu.id ASC`;
+}
+
+function readZcodeNativeUsageMessages(dbPath, sqliteOptions = {}) {
+  let snapshot = null;
+  let effectiveDbPath = dbPath;
+  if (isUncPath(dbPath)) {
+    try {
+      snapshot = snapshotSqliteDb(dbPath);
+      effectiveDbPath = snapshot.path;
+    } catch (_error) {
+      // Preserve the existing direct-read fallback for transient snapshot
+      // failures and UNC implementations that support SQLite locking.
+    }
+  }
+
+  let rows;
+  try {
+    const layout = detectZcodeNativeUsageLayout(effectiveDbPath, sqliteOptions);
+    if (!layout) return null;
+    rows = readSqliteJsonRows(effectiveDbPath, buildZcodeNativeUsageSql(layout), {
+      label: "ZCode",
+      maxBuffer: 50 * 1024 * 1024,
+      timeout: 30_000,
+      ...sqliteOptions,
+      throwOnReadFailure: true,
+    });
+  } catch (_error) {
+    return null;
+  } finally {
+    if (snapshot) snapshot.cleanup();
+  }
+
+  const messages = [];
+  for (const row of rows) {
+    const providerID = String(row?.provider_id || "").trim();
+    const modelID = String(row?.model_id || "").trim();
+    const sessionID = String(row?.session_id || "").trim();
+    const logicalRequestId = String(row?.logical_request_id || "").trim();
+    const attemptIndex = toNonNegativeInt(row?.attempt_index);
+    const id = String(row?.id || "").trim() ||
+      (logicalRequestId ? `${logicalRequestId}#${attemptIndex}` : "");
+    const startedAt = coerceEpochMs(row?.started_at);
+    if (!providerID || !modelID || !sessionID || !id || !startedAt) continue;
+
+    // ZCode's persisted input/output columns are inclusive parents:
+    // cache read/write are subsets of input, and reasoning is a subset of
+    // output. TokenTracker stores those five columns disjointly, so split the
+    // subsets here before entering the shared OpenCode parser. This preserves
+    // sum(input + cache read + cache write + output + reasoning) without
+    // billing either subset twice.
+    const rawInput = toNonNegativeInt(row?.input_tokens);
+    const rawOutput = toNonNegativeInt(row?.output_tokens);
+    const cacheRead = toNonNegativeInt(row?.cache_read_input_tokens);
+    const cacheWrite = toNonNegativeInt(row?.cache_creation_input_tokens);
+    const reasoning = toNonNegativeInt(row?.reasoning_tokens);
+
+    const data = {
+      id,
+      sessionID,
+      role: "assistant",
+      providerID,
+      modelID,
+      time: { created: startedAt, completed: startedAt },
+      tokens: {
+        input: Math.max(0, rawInput - cacheRead - cacheWrite),
+        output: Math.max(0, rawOutput - reasoning),
+        reasoning,
+        cache: {
+          read: cacheRead,
+          write: cacheWrite,
+        },
+      },
+    };
+    if (typeof row?.directory === "string" && row.directory.trim()) {
+      data.path = { cwd: row.directory.trim() };
+    }
+    if (!isZcodeNativeMessage(data)) continue;
+    messages.push({ id, sessionID, timeUpdated: startedAt, data });
+  }
+  return messages;
+}
+
 // Read only genuine ZCode assistant messages (its own GLM models via Z.ai /
 // BigModel), dropping any bundled sub-agent turns. See isZcodeNativeMessage.
 function readZcodeDbMessages(dbPath, sqliteOptions = {}) {
   if (!dbPath || !fssync.existsSync(dbPath)) return [];
+  const nativeMessages = readZcodeNativeUsageMessages(dbPath, sqliteOptions);
+  if (nativeMessages !== null) return nativeMessages;
   const all = readOpencodeDbMessages(dbPath, sqliteOptions);
   return all.filter((m) => isZcodeNativeMessage(m.data));
 }
@@ -4562,7 +4884,8 @@ async function parseOpencodeDbIncremental({
       continue;
     }
 
-    const model = normalizeModelInput(msg?.modelID || msg?.model || msg?.modelId) || DEFAULT_MODEL;
+    const { modelId: dbModelId } = normalizeOpencodeModelFields(msg);
+    const model = dbModelId || DEFAULT_MODEL;
     const bucket = getHourlyBucket(hourlyState, defaultSource, model, bucketStart);
     addTotals(bucket.totals, delta);
     touchedBuckets.add(bucketKey(defaultSource, model, bucketStart));
@@ -16614,36 +16937,55 @@ function isCjkCodePoint(code) {
 // (session transcripts are SQLCipher-encrypted; memory summaries carry no
 // token counts), so the token-count-only queue is intentionally never written
 // for this provider — the cloud hourly table stays clean.
-function resolveTraePath(env = process.env) {
+// Ordered candidate app-dir names. The CN IDE build installs as "Trae CN" on
+// Windows; the international build installs as "TRAE SOLO". Both expose the same
+// iCubeServerData entitlement key, so either is a valid Trae IDE install.
+function traeCandidateAppDirs(env = process.env) {
   const override = env.TOKENTRACKER_TRAE_HOME;
   if (typeof override === "string" && override.trim().length > 0) {
-    return override.trim();
+    return [override.trim()];
   }
   const home = require("node:os").homedir();
   if (process.platform === "darwin") {
-    return path.join(home, "Library", "Application Support", "TRAE SOLO");
+    return [path.join(home, "Library", "Application Support", "TRAE SOLO")];
   }
   if (process.platform === "win32") {
     // A Windows box with no APPDATA still keeps Trae under the standard
     // roaming profile — falling through to the dot-dir below would point at
     // a location Trae never writes.
-    const appData = typeof env.APPDATA === "string" ? env.APPDATA.trim() : "";
-    return appData
-      ? path.join(appData, "TRAE SOLO")
-      : path.join(home, "AppData", "Roaming", "TRAE SOLO");
+    const appData =
+      typeof env.APPDATA === "string" && env.APPDATA.trim()
+        ? env.APPDATA.trim()
+        : path.join(home, "AppData", "Roaming");
+    return [
+      path.join(appData, "TRAE SOLO"),
+      path.join(appData, "Trae CN"),
+    ];
   }
   // Linux and friends: TRAE SOLO ships official builds for macOS/Windows
   // only, so there is no verified app-data layout to default to. Fall back to
   // a deterministic home-dir path (best-effort detection);
   // TOKENTRACKER_TRAE_HOME always wins for unusual installs.
-  return path.join(home, ".trae-solo");
+  return [path.join(home, ".trae-solo")];
+}
+
+// Best-effort: the first candidate that exists on disk (used for display/init
+// detection). When no install is present, returns the primary candidate so
+// callers still get a deterministic path to report.
+function resolveTraePath(env = process.env) {
+  const candidates = traeCandidateAppDirs(env);
+  for (const dir of candidates) {
+    try { if (fssync.existsSync(dir)) return dir; } catch (_error) {}
+  }
+  return candidates.length ? candidates[0] : null;
 }
 
 function resolveTraeStoragePath(env = process.env) {
-  const traHome = resolveTraePath(env);
-  if (!traHome) return null;
-  const p = path.join(traHome, "User", "globalStorage", "storage.json");
-  return fssync.existsSync(p) ? p : null;
+  for (const dir of traeCandidateAppDirs(env)) {
+    const p = path.join(dir, "User", "globalStorage", "storage.json");
+    try { if (fssync.existsSync(p)) return p; } catch (_error) {}
+  }
+  return null;
 }
 
 /**
@@ -17895,6 +18237,7 @@ module.exports = {
   readOpencodeDbMessages,
   readMimoDbMessages,
   readZcodeDbMessages,
+  hasZcodeNativeUsageSchema,
   resolveQoderDbPath,
   resolveQoderDbPaths,
   resolveQoderCnDbPaths,

package/src/lib/trae-cn-config.js

@@ -1,7 +1,7 @@
 "use strict";
 
 /**
- * TRAE Work CN (国内版) config + usage fetch. Phase A: macOS + test-injected
+ * TRAE Work CN (国内版) config + usage fetch. Phase A: macOS + Windows + test-injected
  * paths only. Resolves the CN home, reads the `iCubeAuthInfo://icube.cloudide`
  * blob from User/globalStorage/storage.json (plaintext JSON or Base64 tc v5),
  * and talks to the TRAE CN usage API. JWTs / refresh tokens are never
@@ -85,7 +85,16 @@ function resolveTraeCnHome({ env = process.env, home = os.homedir(), platform =
   if (platform === "darwin") {
     return path.join(home, "Library", "Application Support", TRAE_CN_APP_DIR);
   }
-  // ponytail: this phase is macOS + test-injected paths only.
+  if (platform === "win32") {
+    // Windows keeps the CN Work app under the roaming profile, mirroring the
+    // macOS Application Support layout (APPDATA/TRAE SOLO CN/User/...).
+    const appData =
+      typeof env.APPDATA === "string" && env.APPDATA.trim()
+        ? env.APPDATA.trim()
+        : path.join(home, "AppData", "Roaming");
+    return path.join(appData, TRAE_CN_APP_DIR);
+  }
+  // No verified CN Work app-data layout on other platforms yet.
   return null;
 }