tokentracker-cli 1.4.7 → 1.4.10

package/src/lib/usage-limits.js

@@ -35,6 +35,7 @@ const { fetchDevinLimits } = require("./devin-limits");
 const { fetchQoderLimits, fetchQoderCnLimits } = require("./qoder-limits");
 const { fetchArkCodingPlanLimits } = require("./ark-coding-plan-limits");
 const { fetchArkAgentPlanLimits } = require("./ark-agent-plan-limits");
+const { fetchWorkbuddyLimits, fetchWorkbuddyAiLimits, selectWorkbuddyLimits } = require("./workbuddy-limits");
 const { fetchProviderServiceStatus } = require("./provider-status");
 const { resolveKimiProfile } = require("./kimi-profile");
 const { readSqliteJsonRows, readSqliteJsonRowsAsync } = require("./sqlite-reader");
@@ -56,6 +57,10 @@ const execFileAsync = promisify(cp.execFile);
 const USAGE_LIMITS_SELECTION_KEYS = ["off-on", "off-off", "on-on", "on-off"];
 const cacheBySelection = {};
 for (const key of USAGE_LIMITS_SELECTION_KEYS) cacheBySelection[key] = { data: null, expiresAtMs: 0 };
+// Credential replacements (provider-key save) bump this generation so a fetch
+// that started under the old credential cannot repopulate the memory or disk
+// caches after invalidation.
+let usageLimitsCacheGeneration = 0;
 function usageLimitsSelectionKey(options) {
   const devin = options?.devinEnabled === true ? "on" : "off";
   const antigravity = isAntigravityQuotaDisabled() ? "off" : "on";
@@ -745,13 +750,12 @@ function normalizeCursorUsageSummary(body) {
 }
 
 /**
- * Cursor's Sand "Grok Bot" RPC reports the shared Super weekly pool (same
- * reset and about the same percent as Grok billing's creditUsagePercent).
- * Once the Grok card lists a Grok Bot product share, that Sand bar would
- * show the whole pool a second time under Cursor. Hide it only in that
- * case. A real separate Sand quota (different reset, or a percent more than
- * a point away from the pool) stays. Accounts with no Grok Bot product row
- * keep the Sand bar.
+ * Cursor's Sand RPC reports Grok Bot usage. grok.com shows that as a weekly
+ * Grok Bot limit of its own, separate from the SuperGrok pool and with its own
+ * reset, so the Sand bar normally stays. It is hidden only when it is
+ * indistinguishable from the Grok pool, so the same numbers are not drawn
+ * twice: Grok billing lists a GrokBot product share, and the two windows reset
+ * at the same instant and sit within a point of each other.
  */
 function cursorSandDuplicatesGrokPool(cursor, grok) {
   const sand = cursor?.quaternary_window;
@@ -2603,8 +2607,12 @@ function writeCodexLimitsCache(limits, { home, nowMs = Date.now() } = {}) {
   } catch (_error) {}
 }
 
-function resolveOpencodeGoLimitsCachePath({ home } = {}) {
-  return path.join(home || os.homedir(), ".tokentracker", "tracker", OPENCODE_GO_LIMITS_CACHE_FILE);
+function resolveOpencodeGoLimitsCachePath({ trackerDir, home } = {}) {
+  const dir =
+    typeof trackerDir === "string" && trackerDir.trim()
+      ? trackerDir.trim()
+      : path.join(home || os.homedir(), ".tokentracker", "tracker");
+  return path.join(dir, OPENCODE_GO_LIMITS_CACHE_FILE);
 }
 
 function hasOpencodeGoWindow(limits) {
@@ -2632,10 +2640,15 @@ function normalizeOpencodeGoCachedLimits(
   if (cachedAtMs > nowMs + 60_000) return null;
   if (nowMs - cachedAtMs > effectiveMaxAge) return null;
 
+  const credentialSource =
+    raw?.credential_source === "env" || raw?.credential_source === "stored"
+      ? raw.credential_source
+      : null;
   const cached = {
     configured: true,
     error: null,
     source: typeof raw?.source === "string" ? raw.source : "api",
+    credential_source: credentialSource,
     subscription_status: typeof raw?.subscription_status === "string" ? raw.subscription_status : null,
     plan_label: typeof raw?.plan_label === "string" ? raw.plan_label : null,
     primary_window: isOpencodeGoCacheWindowUsable(raw?.primary_window, { nowMs }) ? raw.primary_window : null,
@@ -2647,8 +2660,8 @@ function normalizeOpencodeGoCachedLimits(
   return hasOpencodeGoWindow(cached) ? cached : null;
 }
 
-function readOpencodeGoLimitsCache({ home, nowMs = Date.now() } = {}) {
-  const cachePath = resolveOpencodeGoLimitsCachePath({ home });
+function readOpencodeGoLimitsCache({ trackerDir, home, nowMs = Date.now() } = {}) {
+  const cachePath = resolveOpencodeGoLimitsCachePath({ trackerDir, home });
   try {
     const parsed = JSON.parse(fs.readFileSync(cachePath, "utf8"));
     return normalizeOpencodeGoCachedLimits(parsed?.opencodeGo, { nowMs });
@@ -2657,12 +2670,17 @@ function readOpencodeGoLimitsCache({ home, nowMs = Date.now() } = {}) {
   }
 }
 
-function writeOpencodeGoLimitsCache(limits, { home, nowMs = Date.now() } = {}) {
+function writeOpencodeGoLimitsCache(limits, { trackerDir, home, nowMs = Date.now() } = {}) {
   if (!limits?.configured || limits.error || !hasOpencodeGoWindow(limits)) return;
-  const cachePath = resolveOpencodeGoLimitsCachePath({ home });
+  const cachePath = resolveOpencodeGoLimitsCachePath({ trackerDir, home });
+  const credentialSource =
+    limits.credential_source === "env" || limits.credential_source === "stored"
+      ? limits.credential_source
+      : null;
   const payload = {
     opencodeGo: {
       source: limits.source || "api",
+      credential_source: credentialSource,
       subscription_status: limits.subscription_status || null,
       plan_label: limits.plan_label || null,
       primary_window: limits.primary_window || null,
@@ -3989,7 +4007,12 @@ async function getUsageLimits(options = {}) {
   if (inFlightBySelection[selection]) {
     return inFlightBySelection[selection];
   }
-  const promise = fetchUsageLimitsUncached({ ...options, selectionKey: selection }).finally(() => {
+  const cacheGeneration = usageLimitsCacheGeneration;
+  const promise = fetchUsageLimitsUncached({
+    ...options,
+    selectionKey: selection,
+    cacheGeneration,
+  }).finally(() => {
     if (inFlightBySelection[selection] === promise) inFlightBySelection[selection] = null;
   });
   inFlightBySelection[selection] = promise;
@@ -4002,6 +4025,7 @@ async function getUsageLimits(options = {}) {
  */
 async function fetchUsageLimitsUncached({
   home,
+  trackerDir,
   env,
   platform,
   securityRunner,
@@ -4013,6 +4037,10 @@ async function fetchUsageLimitsUncached({
   providerTimeoutMs = DEFAULT_PROVIDER_TIMEOUT_MS,
   forceRefresh = false,
   devinEnabled = false,
+  // Generation captured by getUsageLimits before the fetch began. A provider
+  // key save bumps it, and this fetch must then stop at returning its own
+  // result: publishing it would restore pre-save data for every later caller.
+  cacheGeneration = usageLimitsCacheGeneration,
   // Selection key fixed by getUsageLimits before the fetch began. Re-deriving it
   // here would let an opt-out flipped mid-fetch store a pre-opt-out result in the
   // post-opt-out slot.
@@ -4079,7 +4107,7 @@ async function fetchUsageLimitsUncached({
     : null;
 
   const providerFetch = withFetchTimeout(fetchImpl, providerTimeoutMs);
-  const [claudeResult, codexResult, cursor, kimi, gemini, kiro, antigravity, copilot, grok, zcode, opencodeGoRaw, qoder, qoderCn, codingPlan, agentPlan, commandCodeRaw, devinRaw, claudeServiceStatus] = await Promise.all([
+  const [claudeResult, codexResult, cursor, kimi, gemini, kiro, antigravity, copilot, grok, zcode, opencodeGoRaw, qoder, qoderCn, codingPlan, agentPlan, commandCodeRaw, devinRaw, workbuddy, workbuddyAi, claudeServiceStatus] = await Promise.all([
     claudeToken && !freshClaudeCache && !claudeRetryAtMs
       ? withProviderTimeout(fetchClaudeUsageLimits(claudeToken, { fetchImpl: providerFetch, maxAttempts: 1 }), "Claude", providerTimeoutMs).then(
           (value) => ({ status: "fulfilled", value }),
@@ -4131,7 +4159,7 @@ async function fetchUsageLimitsUncached({
     // OpenCode Go: authoritative subscription windows come from the dashboard
     // scrape; local opencode.db cost is available only as an explicit estimate.
     // See src/lib/opencode-go-limits.js.
-    withProviderTimeout(fetchOpencodeGoLimits({ home, env, fetchImpl: providerFetch }), "OpenCode Go", providerTimeoutMs)
+    withProviderTimeout(fetchOpencodeGoLimits({ home, trackerDir, env, fetchImpl: providerFetch }), "OpenCode Go", providerTimeoutMs)
       .catch((reason) => ({ configured: true, error: reason?.message || "Unknown error" })),
     withProviderTimeout(
       fetchQoderLimits({
@@ -4189,7 +4217,7 @@ async function fetchUsageLimitsUncached({
     // outer race bounds the whole slot: without it three sequential per-request
     // timeouts (whoami, then credits+subscriptions) could hold the aggregate
     // ~2x longer than any sibling provider.
-    withProviderTimeout(fetchCommandcodeLimits({ home, env, fetchImpl: providerFetch }), "CommandCode", providerTimeoutMs)
+    withProviderTimeout(fetchCommandcodeLimits({ home, trackerDir, env, fetchImpl: providerFetch }), "CommandCode", providerTimeoutMs)
       .then(
         (value) => ({ status: "fulfilled", value }),
         (reason) => ({ status: "rejected", reason }),
@@ -4204,6 +4232,13 @@ async function fetchUsageLimitsUncached({
         (value) => ({ status: "fulfilled", value }),
         (reason) => ({ status: "rejected", reason }),
       ),
+    // WorkBuddy CN and WorkBuddy AI are separate editions. Each one calls only
+    // its own local authenticated billing proxy. Abort closes the IPC socket
+    // when the provider budget expires. Reward credits are not a plan window.
+    withAbortableProviderTimeout((signal) => fetchWorkbuddyLimits({ home, env, fetchImpl: providerFetch, signal }), "WorkBuddy", providerTimeoutMs)
+      .catch((reason) => ({ configured: true, error: reason?.message || "Unknown error" })),
+    withAbortableProviderTimeout((signal) => fetchWorkbuddyAiLimits({ home, env, fetchImpl: providerFetch, signal }), "WorkBuddy AI", providerTimeoutMs)
+      .catch((reason) => ({ configured: true, error: reason?.message || "Unknown error" })),
     // Public status-page probe (fail-soft, own 5-min cache in provider-status.js).
     // Only probed for configured accounts — without a token the Claude section
     // never renders, so the reading would have nowhere to go.
@@ -4382,9 +4417,11 @@ async function fetchUsageLimitsUncached({
       stale: false,
       cached_at: new Date(nowMs).toISOString(),
     };
-    writeOpencodeGoLimitsCache(opencodeGo, { home, nowMs });
+    if (cacheGeneration === usageLimitsCacheGeneration) {
+      writeOpencodeGoLimitsCache(opencodeGo, { trackerDir, home, nowMs });
+    }
   } else {
-    const cached = readOpencodeGoLimitsCache({ home, nowMs });
+    const cached = readOpencodeGoLimitsCache({ trackerDir, home, nowMs });
     if (cached) {
       opencodeGo = cached;
     } else {
@@ -4468,6 +4505,9 @@ async function fetchUsageLimitsUncached({
     qoderCn: withPlanLabel(qoderCn, qoderCn?.plan_label, "Qoder CN"),
     codingPlan: withPlanLabel(codingPlan, codingPlan?.plan_label, "Ark Coding Plan"),
     agentPlan: withPlanLabel(agentPlan, agentPlan?.plan_label, "Ark Agent Plan"),
+    // WorkBuddy preserves provider English names and official edition labels
+    // verbatim (including Free, STANDARD and Pro-Trial).
+    workbuddy: selectWorkbuddyLimits(workbuddy, workbuddyAi),
   };
 
   for (const [providerName, provider] of Object.entries(data)) {
@@ -4491,16 +4531,22 @@ async function fetchUsageLimitsUncached({
     };
   }
 
-  cacheBySelection[selectionKey || usageLimitsSelectionKey({ devinEnabled })] = {
-    data,
-    expiresAtMs: cacheExpiresAtMs(data, nowMs),
-  };
+  if (cacheGeneration === usageLimitsCacheGeneration) {
+    cacheBySelection[selectionKey || usageLimitsSelectionKey({ devinEnabled })] = {
+      data,
+      expiresAtMs: cacheExpiresAtMs(data, nowMs),
+    };
+  }
   return data;
 }
 
-function resetUsageLimitsCache() {
+function resetUsageLimitsCache({ invalidateInFlight = false } = {}) {
+  if (invalidateInFlight) usageLimitsCacheGeneration += 1;
   for (const key of USAGE_LIMITS_SELECTION_KEYS) {
     cacheBySelection[key] = { data: null, expiresAtMs: 0 };
+    // Ordinary resets still let a running fetch serve refresh=1 callers; a
+    // credential replacement must not, so drop the old-key promises too.
+    if (invalidateInFlight) inFlightBySelection[key] = null;
   }
 }
 
@@ -4508,6 +4554,8 @@ module.exports = {
   getUsageLimits,
   normalizePlanLabel,
   resetUsageLimitsCache,
+  OPENCODE_GO_LIMITS_CACHE_FILE,
+  resolveOpencodeGoLimitsCachePath,
   cacheExpiresAtMs,
   runCommand,
   extractGeminiOauthClientCredentials,
@@ -4549,4 +4597,6 @@ module.exports = {
   fetchCommandcodeLimits,
   fetchQoderLimits,
   fetchQoderCnLimits,
+  fetchWorkbuddyLimits,
+  fetchWorkbuddyAiLimits,
 };

package/src/lib/workbuddy-ipc.js

@@ -0,0 +1,191 @@
+"use strict";
+
+// WorkBuddy's authenticated local proxy. The ticket proves both peers; the
+// desktop daemon adds its own access token to relative billing requests.
+const fs = require("node:fs");
+const path = require("node:path");
+const net = require("node:net");
+const crypto = require("node:crypto");
+
+const UNAVAILABLE = "Start WorkBuddy and sign in to read its subscription quota.";
+const MAX_FRAME_BYTES = 1024 * 1024;
+const BILLING_PATHS = new Set([
+  "/billing/meter/get-user-resource-summary",
+  "/billing/meter/get-user-resource-paid-packages",
+  "/billing/meter/get-user-resource-free-packages",
+]);
+
+function readWorkbuddyIpcDiscovery(workbuddyHome) {
+  if (!workbuddyHome) return null;
+  let fd;
+  try {
+    fd = fs.openSync(path.join(workbuddyHome, "wbipc", "endpoint.json"),
+      fs.constants.O_RDONLY | (fs.constants.O_NOFOLLOW || 0));
+    const stat = fs.fstatSync(fd);
+    if (!stat.isFile() || stat.size > 16384 || (process.platform !== "win32"
+      && (stat.uid !== process.getuid() || (stat.mode & 0o077) !== 0))) {
+      throw new Error("Untrusted discovery file.");
+    }
+    const value = JSON.parse(fs.readFileSync(fd, "utf8"));
+    if (typeof value.ticket !== "string" || !/^[\w-]{43}$/.test(value.ticket)
+      || typeof value.endpoint !== "string") throw new Error("Invalid discovery file.");
+    if (process.platform === "win32") {
+      if (!/^\\\\\.\\pipe\\wbipc-[\w-]+$/.test(value.endpoint)) throw new Error("Invalid pipe.");
+    } else {
+      if (!path.isAbsolute(value.endpoint)) throw new Error("Invalid socket.");
+      const socket = fs.lstatSync(value.endpoint);
+      if (!socket.isSocket() || socket.uid !== process.getuid() || (socket.mode & 0o077) !== 0) {
+        throw new Error("Untrusted socket.");
+      }
+    }
+    return { endpoint: value.endpoint, ticket: value.ticket };
+  } catch (error) {
+    if (error.code === "ENOENT" || error.code === "ENOTDIR") return null;
+    // Never expose discovery contents, tickets, paths or JSON parser excerpts.
+    throw new Error("WorkBuddy local connection information is unavailable.");
+  } finally {
+    if (fd !== undefined) fs.closeSync(fd);
+  }
+}
+
+function computeProof(ticket, role, endpoint, clientNonce, serverNonce) {
+  const parts = [];
+  for (const field of [role === "server" ? "wbipc-s" : "wbipc-c", "1", endpoint, clientNonce, serverNonce]) {
+    const bytes = Buffer.from(field, "utf8");
+    const length = Buffer.alloc(4);
+    length.writeUInt32BE(bytes.length);
+    parts.push(length, bytes);
+  }
+  return crypto.createHmac("sha256", ticket).update(Buffer.concat(parts)).digest("base64url");
+}
+
+function openWorkbuddyIpc({ discovery, signal, timeoutMs = 12000, connectImpl = net.createConnection }) {
+  return new Promise((resolve, reject) => {
+    const { endpoint, ticket } = discovery;
+    const clientNonce = crypto.randomBytes(16).toString("base64url");
+    const pending = new Map();
+    let socket;
+    let phase = "challenge";
+    let channel;
+    let nextId = 1;
+    let buffer = Buffer.alloc(0);
+    let closed = false;
+    const timer = setTimeout(() => fail(new Error("WorkBuddy quota request timed out.")), timeoutMs);
+    const onAbort = () => fail(new Error("WorkBuddy quota request was cancelled."));
+
+    function fail(error) {
+      if (closed) return;
+      closed = true;
+      clearTimeout(timer);
+      signal?.removeEventListener("abort", onAbort);
+      reject(error);
+      for (const handlers of pending.values()) handlers.reject(error);
+      pending.clear();
+      socket?.destroy();
+    }
+    function send(packet) {
+      socket.write(JSON.stringify(packet) + "\n");
+    }
+    function rpc(method, params, mode) {
+      if (closed) return Promise.reject(new Error(UNAVAILABLE));
+      return new Promise((resolve, reject) => {
+        const id = nextId++;
+        pending.set(id, { resolve, reject });
+        send({ jsonrpc: "2.0", id, method, params, ...(mode ? { mode } : {}) });
+      });
+    }
+    const client = {
+      close: () => fail(new Error("WorkBuddy local connection closed.")),
+      async postBilling(apiPath, body) {
+        if (!BILLING_PATHS.has(apiPath)) throw new Error("Unsupported WorkBuddy quota request.");
+        const response = await rpc(`${channel}/http.fetch`, {
+          method: "POST",
+          path: apiPath,
+          headers: { "content-type": "application/json", accept: "application/json" },
+          body_b64: Buffer.from(JSON.stringify(body)).toString("base64"),
+        }, "call");
+        if (response?.status === 401 || response?.status === 403) {
+          const error = new Error("WorkBuddy sign-in expired.");
+          error.code = "AUTH_EXPIRED";
+          throw error;
+        }
+        if (response?.status !== 200) throw new Error("WorkBuddy quota request failed.");
+        try {
+          if (typeof response.body_b64 !== "string" || !/^[A-Za-z0-9+/]*={0,2}$/.test(response.body_b64)) {
+            throw new Error("Invalid body.");
+          }
+          return JSON.parse(Buffer.from(response.body_b64, "base64").toString("utf8"));
+        } catch {
+          throw new Error("WorkBuddy quota response was malformed.");
+        }
+      },
+    };
+    function onPacket(packet) {
+      if (!packet || typeof packet !== "object") throw new Error("Invalid packet.");
+      if (phase === "challenge" && packet.type === "session_challenge" && packet.protocol === 1) {
+        if (typeof packet.server_nonce !== "string" || !/^[\w-]{16,128}$/.test(packet.server_nonce)
+          || typeof packet.server_proof !== "string") throw new Error("Invalid challenge.");
+        const expected = Buffer.from(computeProof(ticket, "server", endpoint, clientNonce, packet.server_nonce));
+        const actual = Buffer.from(packet.server_proof);
+        if (expected.length !== actual.length || !crypto.timingSafeEqual(expected, actual)) {
+          throw new Error("Invalid server proof.");
+        }
+        phase = "ack";
+        send({ type: "session_prove", client_proof: computeProof(ticket, "client", endpoint, clientNonce, packet.server_nonce) });
+      } else if (phase === "ack" && packet.type === "session_hello_ack" && packet.protocol === 1
+        && Array.isArray(packet.pipes) && packet.pipes.includes("wb.request")) {
+        phase = "ready";
+        rpc("broker/GetPipe", { pipe: "wb.request" }).then((result) => {
+          if (closed) return;
+          if (result?.channel !== "c:wb.request" || !Array.isArray(result.methods)
+            || !result.methods.includes("http.fetch")) {
+            fail(new Error("WorkBuddy billing proxy is unavailable."));
+            return;
+          }
+          channel = result.channel;
+          resolve(client);
+        }, () => fail(new Error(UNAVAILABLE)));
+      } else if (phase === "ready" && packet.jsonrpc === "2.0" && pending.has(packet.id)) {
+        const handlers = pending.get(packet.id);
+        pending.delete(packet.id);
+        if (packet.error) handlers.reject(new Error("WorkBuddy local quota request was rejected."));
+        else handlers.resolve(packet.result);
+      } else {
+        throw new Error("Unexpected packet.");
+      }
+    }
+    signal?.addEventListener("abort", onAbort, { once: true });
+    if (signal?.aborted) { onAbort(); return; }
+    try {
+      socket = connectImpl(endpoint);
+      socket.on("connect", () => send({
+        type: "session_hello", protocol_min: 1, protocol_max: 1,
+        client_nonce: clientNonce,
+        ticket_id: crypto.createHash("sha256").update(ticket, "utf8").digest("hex").slice(0, 16),
+        client: { kind: "tokentracker", id: "subscription-quota", version: "1" },
+      }));
+      socket.on("data", (chunk) => {
+        if (closed) return;
+        buffer = Buffer.concat([buffer, chunk]);
+        if (buffer.length > MAX_FRAME_BYTES) { fail(new Error("WorkBuddy quota response was too large.")); return; }
+        try {
+          let index;
+          while ((index = buffer.indexOf(10)) >= 0 && !closed) {
+            const line = buffer.subarray(0, index);
+            buffer = buffer.subarray(index + 1);
+            if (line.length) onPacket(JSON.parse(line.toString("utf8")));
+          }
+        } catch {
+          fail(new Error("WorkBuddy local connection authentication or protocol failed."));
+        }
+      });
+      socket.on("error", () => fail(new Error(UNAVAILABLE)));
+      socket.on("close", () => fail(new Error(UNAVAILABLE)));
+      socket.on("end", () => fail(new Error(UNAVAILABLE)));
+    } catch {
+      fail(new Error(UNAVAILABLE));
+    }
+  });
+}
+
+module.exports = { UNAVAILABLE, readWorkbuddyIpcDiscovery, openWorkbuddyIpc };