From 04edf7d9c4e2e2c1ff80b72fb046377bdfb328fa Mon Sep 17 00:00:00 2001
From: Haileyesus <118998054+blackmammoth@users.noreply.github.com>
Date: Fri, 29 May 2026 17:40:44 +0300
Subject: [PATCH] fix: harden token usage reporting

Token usage is shown as concrete counts, so malformed provider payloads must not leak NaN.

Gemini can emit token fields as strings or invalid values, so non-finite values now fall back to 0.

OpenCode token reads happen while the CLI is shutting down, when the DB may be missing or locked.

Those failures now return null instead of interrupting session completion.

/cost no longer invents an input breakdown from an aggregate total.

When a provider only supplies total usage, the UI now says the breakdown is unavailable.

This keeps the display honest instead of presenting made-up input and output rows.

Verification: npm run typecheck; targeted eslint.
---
 server/gemini-response-handler.js             | 11 +++++---
 server/opencode-cli.js                        |  9 +++++--
 server/routes/commands.js                     | 18 ++++++-------
 .../view/subcomponents/CommandResultModal.tsx | 27 ++++++++++++++++---
 4 files changed, 45 insertions(+), 20 deletions(-)

diff --git a/server/gemini-response-handler.js b/server/gemini-response-handler.js
index f57ad072..6f314443 100644
--- a/server/gemini-response-handler.js
+++ b/server/gemini-response-handler.js
@@ -7,10 +7,13 @@ function buildGeminiTokenBudget(tokens) {
     return null;
   }
 
-  const inputTokens = Number(tokens.input || 0);
-  const outputTokens = Number(tokens.output || 0);
-  const used = Number(tokens.total || inputTokens + outputTokens || 0);
-  if (used <= 0) {
+  const parsedInputTokens = Number(tokens.input);
+  const parsedOutputTokens = Number(tokens.output);
+  const inputTokens = Number.isFinite(parsedInputTokens) ? parsedInputTokens : 0;
+  const outputTokens = Number.isFinite(parsedOutputTokens) ? parsedOutputTokens : 0;
+  const parsedUsed = Number(tokens.total);
+  const used = Number.isFinite(parsedUsed) ? parsedUsed : inputTokens + outputTokens;
+  if (!Number.isFinite(used) || used <= 0) {
     return null;
   }
 
diff --git a/server/opencode-cli.js b/server/opencode-cli.js
index 119c5ee4..6f3e0f65 100644
--- a/server/opencode-cli.js
+++ b/server/opencode-cli.js
@@ -28,8 +28,9 @@ function readOpenCodeTokenUsage(sessionId) {
     return null;
   }
 
-  const db = new Database(dbPath, { readonly: true, fileMustExist: true });
+  let db = null;
   try {
+    db = new Database(dbPath, { readonly: true, fileMustExist: true });
     const columns = db.prepare('PRAGMA table_info(session)').all();
     const columnNames = new Set(columns.map((column) => column.name));
     const requiredColumns = ['tokens_input', 'tokens_output', 'tokens_reasoning', 'tokens_cache_read', 'tokens_cache_write'];
@@ -72,8 +73,12 @@ function readOpenCodeTokenUsage(sessionId) {
         output: outputTokens,
       },
     };
+  } catch {
+    return null;
   } finally {
-    db.close();
+    if (db) {
+      db.close();
+    }
   }
 }
 
diff --git a/server/routes/commands.js b/server/routes/commands.js
index 82a0e7e0..09022e0c 100644
--- a/server/routes/commands.js
+++ b/server/routes/commands.js
@@ -291,12 +291,6 @@ Custom commands can be created in:
     const hasTokenBreakdown = inputTokensRaw > 0 || outputTokens > 0;
     const used = reportedUsed || inputTokensRaw + outputTokens;
 
-    // If we only have total used tokens, keep the list populated without guessing output.
-    const inputTokens =
-      hasTokenBreakdown
-        ? inputTokensRaw
-        : used;
-
     return {
       type: "builtin",
       action: "cost",
@@ -305,10 +299,14 @@ Custom commands can be created in:
           used,
           total,
         },
-        tokenBreakdown: {
-          input: inputTokens,
-          output: outputTokens,
-        },
+        ...(hasTokenBreakdown
+          ? {
+              tokenBreakdown: {
+                input: inputTokensRaw,
+                output: outputTokens,
+              },
+            }
+          : {}),
         provider,
         model,
       },
diff --git a/src/components/chat/view/subcomponents/CommandResultModal.tsx b/src/components/chat/view/subcomponents/CommandResultModal.tsx
index d51f7c96..2e391ebe 100644
--- a/src/components/chat/view/subcomponents/CommandResultModal.tsx
+++ b/src/components/chat/view/subcomponents/CommandResultModal.tsx
@@ -495,12 +495,31 @@ function CostContent({ data }: { data: CostCommandData }) {
   const total = Number(data.tokenUsage?.total ?? 0);
   const model = data.model || 'Unknown';
   const provider = getProviderLabel(data.provider, data.provider || 'Unknown');
-  const inputTokens = Number(data.tokenBreakdown?.input ?? 0);
-  const outputTokens = Number(data.tokenBreakdown?.output ?? 0);
+  const hasBreakdown =
+    typeof data.tokenBreakdown?.input === 'number' ||
+    typeof data.tokenBreakdown?.output === 'number';
   const usageRows = [
     { label: 'Total tokens used', value: formatNumber(used), icon: Activity },
-    { label: 'Input tokens', value: formatNumber(inputTokens), icon: TerminalSquare },
-    { label: 'Output tokens', value: formatNumber(outputTokens), icon: Coins },
+    ...(hasBreakdown
+      ? [
+          {
+            label: 'Input tokens',
+            value: formatNumber(Number(data.tokenBreakdown?.input ?? 0)),
+            icon: TerminalSquare,
+          },
+          {
+            label: 'Output tokens',
+            value: formatNumber(Number(data.tokenBreakdown?.output ?? 0)),
+            icon: Coins,
+          },
+        ]
+      : [
+          {
+            label: 'Breakdown',
+            value: 'Unavailable',
+            icon: TerminalSquare,
+          },
+        ]),
     ...(total > 0
       ? [{ label: 'Context window', value: formatNumber(total), icon: Gauge }]
       : []),