From 00f8d113df566fafb606bdbea8edc38d23f263d1 Mon Sep 17 00:00:00 2001
From: Ismail Pelaseyed <homanp@gmail.com>
Date: Wed, 23 Jul 2025 14:54:47 +0200
Subject: [PATCH] fix issues with token counting (#33)

---
 package.json                          |  2 +-
 src/agent/grok-agent.ts               | 15 +++++++++++----
 src/ui/components/loading-spinner.tsx |  3 ++-
 src/utils/token-counter.ts            | 17 +++++++++++++++++
 4 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/package.json b/package.json
index 018b867..f455f80 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@vibe-kit/grok-cli",
-  "version": "0.0.5",
+  "version": "0.0.6",
   "description": "An open-source AI agent that brings the power of Grok directly into your terminal.",
   "main": "dist/index.js",
   "bin": {
diff --git a/src/agent/grok-agent.ts b/src/agent/grok-agent.ts
index 6abf627..2f98a50 100644
--- a/src/agent/grok-agent.ts
+++ b/src/agent/grok-agent.ts
@@ -325,7 +325,7 @@ Current working directory: ${process.cwd()}`,
     this.messages.push({ role: "user", content: message });
 
     // Calculate input tokens
-    const inputTokens = this.tokenCounter.countMessageTokens(
+    let inputTokens = this.tokenCounter.countMessageTokens(
       this.messages as any
     );
     yield {
@@ -396,9 +396,9 @@ Current working directory: ${process.cwd()}`,
           if (chunk.choices[0].delta?.content) {
             accumulatedContent += chunk.choices[0].delta.content;
 
-            // Update token count in real-time
-            const currentOutputTokens =
-              this.tokenCounter.estimateStreamingTokens(accumulatedContent);
+            // Update token count in real-time including accumulated content and any tool calls
+            const currentOutputTokens = this.tokenCounter.estimateStreamingTokens(accumulatedContent) +
+              (accumulatedMessage.tool_calls ? this.tokenCounter.countTokens(JSON.stringify(accumulatedMessage.tool_calls)) : 0);
             totalOutputTokens = currentOutputTokens;
 
             yield {
@@ -483,6 +483,13 @@ Current working directory: ${process.cwd()}`,
             });
           }
 
+          // Update token count after processing all tool calls to include tool results
+          inputTokens = this.tokenCounter.countMessageTokens(this.messages as any);
+          yield {
+            type: "token_count",
+            tokenCount: inputTokens + totalOutputTokens,
+          };
+
           // Continue the loop to get the next response (which might have more tool calls)
         } else {
           // No tool calls, we're done
diff --git a/src/ui/components/loading-spinner.tsx b/src/ui/components/loading-spinner.tsx
index f36d562..8baa804 100644
--- a/src/ui/components/loading-spinner.tsx
+++ b/src/ui/components/loading-spinner.tsx
@@ -1,5 +1,6 @@
 import React, { useState, useEffect } from "react";
 import { Box, Text } from "ink";
+import { formatTokenCount } from "../../utils/token-counter";
 
 interface LoadingSpinnerProps {
   isActive: boolean;
@@ -61,7 +62,7 @@ export function LoadingSpinner({ isActive, processingTime, tokenCount }: Loading
         {spinnerFrames[spinnerFrame]} {loadingTexts[loadingTextIndex]}{" "}
       </Text>
       <Text color="gray">
-        ({processingTime}s · ↑ {tokenCount} tokens · esc to interrupt)
+        ({processingTime}s · ↑ {formatTokenCount(tokenCount)} tokens · esc to interrupt)
       </Text>
     </Box>
   );
diff --git a/src/utils/token-counter.ts b/src/utils/token-counter.ts
index a049d5a..9404a02 100644
--- a/src/utils/token-counter.ts
+++ b/src/utils/token-counter.ts
@@ -66,6 +66,23 @@ export class TokenCounter {
   }
 }
 
+/**
+ * Format token count for display (e.g., 1.2k for 1200)
+ */
+export function formatTokenCount(count: number): string {
+  if (count <= 999) {
+    return count.toString();
+  }
+  
+  if (count < 1_000_000) {
+    const k = count / 1000;
+    return k % 1 === 0 ? `${k}k` : `${k.toFixed(1)}k`;
+  }
+  
+  const m = count / 1_000_000;
+  return m % 1 === 0 ? `${m}m` : `${m.toFixed(1)}m`;
+}
+
 /**
  * Create a token counter instance
  */