Show token bar at bottom of chat input (#33)

src/components/chat/ChatInput.tsx

@@ -15,6 +15,7 @@ import {
   Database,
   ChevronsUpDown,
   ChevronsDownUp,
+  BarChart2,
 } from "lucide-react";
 import type React from "react";
 import { useCallback, useEffect, useRef, useState } from "react";
@@ -22,7 +23,7 @@ import { ModelPicker } from "@/components/ModelPicker";
 import { useSettings } from "@/hooks/useSettings";
 import { IpcClient } from "@/ipc/ipc_client";
 import { chatInputValueAtom, chatMessagesAtom } from "@/atoms/chatAtoms";
-import { useAtom, useSetAtom } from "jotai";
+import { atom, useAtom, useSetAtom } from "jotai";
 import { useStreamChat } from "@/hooks/useStreamChat";
 import { useChats } from "@/hooks/useChats";
 import { selectedAppIdAtom } from "@/atoms/appAtoms";
@@ -46,6 +47,9 @@ import { useRunApp } from "@/hooks/useRunApp";
 import { AutoApproveSwitch } from "../AutoApproveSwitch";
 import { usePostHog } from "posthog-js/react";
 import { CodeHighlight } from "./CodeHighlight";
+import { TokenBar } from "./TokenBar";
+
+const showTokenBarAtom = atom(false);
 
 export function ChatInput({ chatId }: { chatId?: number }) {
   const posthog = usePostHog();
@@ -60,6 +64,7 @@ export function ChatInput({ chatId }: { chatId?: number }) {
   const [isRejecting, setIsRejecting] = useState(false); // State for rejecting
   const [messages, setMessages] = useAtom<Message[]>(chatMessagesAtom);
   const setIsPreviewOpen = useSetAtom(isPreviewOpenAtom);
+  const [showTokenBar, setShowTokenBar] = useAtom(showTokenBarAtom);
 
   const { refreshAppIframe } = useRunApp();
 
@@ -274,14 +279,26 @@ export function ChatInput({ chatId }: { chatId?: number }) {
               </button>
             )}
           </div>
-          <div className="px-2 pb-2">
-            <ModelPicker
-              selectedModel={settings.selectedModel}
-              onModelSelect={(model) =>
-                updateSettings({ selectedModel: model })
-              }
-            />
+          <div className="pl-2 pr-1 flex items-center justify-between">
+            <div className="pb-2">
+              <ModelPicker
+                selectedModel={settings.selectedModel}
+                onModelSelect={(model) =>
+                  updateSettings({ selectedModel: model })
+                }
+              />
+            </div>
+            <button
+              onClick={() => setShowTokenBar(!showTokenBar)}
+              className="flex items-center px-2 py-1 text-xs text-muted-foreground hover:bg-muted rounded"
+              title={showTokenBar ? "Hide token usage" : "Show token usage"}
+            >
+              <BarChart2 size={14} className="mr-1" />
+              {showTokenBar ? "Hide tokens" : "Tokens"}
+            </button>
           </div>
+          {/* TokenBar is only displayed when showTokenBar is true */}
+          {showTokenBar && <TokenBar chatId={chatId} />}
         </div>
       </div>
     </>

src/components/chat/TokenBar.tsx

@@ -0,0 +1,130 @@
+import React, { useEffect, useState } from "react";
+import {
+  Tooltip,
+  TooltipContent,
+  TooltipProvider,
+  TooltipTrigger,
+} from "@/components/ui/tooltip";
+import { useCountTokens } from "@/hooks/useCountTokens";
+import { MessageSquare, Code, Bot, AlignLeft } from "lucide-react";
+import { chatInputValueAtom } from "@/atoms/chatAtoms";
+import { useAtom } from "jotai";
+import { useSettings } from "@/hooks/useSettings";
+
+interface TokenBarProps {
+  chatId?: number;
+}
+
+export function TokenBar({ chatId }: TokenBarProps) {
+  const [inputValue] = useAtom(chatInputValueAtom);
+  const { countTokens, result } = useCountTokens();
+  const [error, setError] = useState<string | null>(null);
+  const { settings } = useSettings();
+  useEffect(() => {
+    if (!chatId) return;
+    // Mark this as used, we need to re-trigger token count
+    // when selected model changes.
+    void settings?.selectedModel;
+
+    const debounceTimer = setTimeout(() => {
+      countTokens(chatId, inputValue).catch((err) => {
+        setError("Failed to count tokens");
+        console.error("Token counting error:", err);
+      });
+    }, 500);
+
+    return () => clearTimeout(debounceTimer);
+  }, [chatId, inputValue, countTokens, settings?.selectedModel]);
+
+  if (!chatId || !result) {
+    return null;
+  }
+
+  const {
+    totalTokens,
+    messageHistoryTokens,
+    codebaseTokens,
+    systemPromptTokens,
+    inputTokens,
+    contextWindow,
+  } = result;
+
+  const percentUsed = Math.min((totalTokens / contextWindow) * 100, 100);
+
+  // Calculate widths for each token type
+  const messageHistoryPercent = (messageHistoryTokens / contextWindow) * 100;
+  const codebasePercent = (codebaseTokens / contextWindow) * 100;
+  const systemPromptPercent = (systemPromptTokens / contextWindow) * 100;
+  const inputPercent = (inputTokens / contextWindow) * 100;
+
+  return (
+    <div className="px-4 pb-2 text-xs">
+      <TooltipProvider>
+        <Tooltip>
+          <TooltipTrigger asChild>
+            <div className="w-full">
+              <div className="flex justify-between mb-1 text-xs text-muted-foreground">
+                <span>Tokens: {totalTokens.toLocaleString()}</span>
+                <span>
+                  {Math.round(percentUsed)}% of{" "}
+                  {(contextWindow / 1000).toFixed(0)}K
+                </span>
+              </div>
+              <div className="w-full h-1.5 bg-muted rounded-full overflow-hidden flex">
+                {/* Message history tokens */}
+                <div
+                  className="h-full bg-blue-400"
+                  style={{ width: `${messageHistoryPercent}%` }}
+                />
+                {/* Codebase tokens */}
+                <div
+                  className="h-full bg-green-400"
+                  style={{ width: `${codebasePercent}%` }}
+                />
+                {/* System prompt tokens */}
+                <div
+                  className="h-full bg-purple-400"
+                  style={{ width: `${systemPromptPercent}%` }}
+                />
+                {/* Input tokens */}
+                <div
+                  className="h-full bg-yellow-400"
+                  style={{ width: `${inputPercent}%` }}
+                />
+              </div>
+            </div>
+          </TooltipTrigger>
+          <TooltipContent side="top" className="w-64 p-2">
+            <div className="space-y-1">
+              <div className="font-medium">Token Usage Breakdown</div>
+              <div className="grid grid-cols-[20px_1fr_auto] gap-x-2 items-center">
+                <MessageSquare size={12} className="text-blue-500" />
+                <span>Message History</span>
+                <span>{messageHistoryTokens.toLocaleString()}</span>
+
+                <Code size={12} className="text-green-500" />
+                <span>Codebase</span>
+                <span>{codebaseTokens.toLocaleString()}</span>
+
+                <Bot size={12} className="text-purple-500" />
+                <span>System Prompt</span>
+                <span>{systemPromptTokens.toLocaleString()}</span>
+
+                <AlignLeft size={12} className="text-yellow-500" />
+                <span>Current Input</span>
+                <span>{inputTokens.toLocaleString()}</span>
+              </div>
+              <div className="pt-1 border-t border-border">
+                <div className="flex justify-between font-medium">
+                  <span>Total</span>
+                  <span>{totalTokens.toLocaleString()}</span>
+                </div>
+              </div>
+            </div>
+          </TooltipContent>
+        </Tooltip>
+      </TooltipProvider>
+      {error && <div className="text-red-500 text-xs mt-1">{error}</div>}
+    </div>
+  );
+}