feat: add token usage tracking and display in ChatInterface

- Introduced TokenCountInfo component to display total token usage in the chat interface.
- Updated EphemeralStore to manage session-specific token usage data.
- Enhanced event handling to capture and store token usage from incoming events.
- Integrated token usage display within the ChatInterface for improved user feedback.

Author: milisp

src/components/chat/ChatInterface.tsx

@@ -17,6 +17,7 @@ import { ForkOriginBanner } from './ForkOriginBanner';
 import { useEphemeralStore } from '@/stores/EphemeralStore';
 import { ChangesSummary } from './ChangesSummary';
 import { ModelSelector } from "./ModelSelector";
+import TokenCountInfo from "@/components/common/TokenCountInfo";
 
 interface ChatInterfaceProps {
   sessionId: string;
@@ -107,6 +108,7 @@ export const ChatInterface: React.FC<ChatInterfaceProps> = ({
   const messages = [...sessionMessages];
   const isLoading = currentConversation?.isLoading || false;
   const fileDiffMap = useEphemeralStore((s) => s.sessionFileDiffs[activeSessionId]);
+  const sessionUsage = useEphemeralStore((s) => s.sessionTokenUsage[activeSessionId]);
 
   // Update activeSessionId when sessionId prop changes
   useEffect(() => {
@@ -480,7 +482,8 @@ export const ChatInterface: React.FC<ChatInterfaceProps> = ({
           />
           <ModelSelector />
           <ReasoningEffortSelector />
-          </div>
+          <TokenCountInfo usage={sessionUsage} />
+        </div>
       </div>
     </div>
   );

src/components/common/TokenCountInfo.tsx

@@ -0,0 +1,30 @@
+type TokenCountShape = {
+  input_tokens?: number;
+  cached_input_tokens?: number;
+  output_tokens?: number;
+  reasoning_output_tokens?: number;
+  total_tokens: number;
+};
+
+interface TokenCountInfoProps {
+  usage: TokenCountShape | null | undefined;
+  className?: string;
+}
+
+export function TokenCountInfo({ usage, className }: TokenCountInfoProps) {
+  if (!usage || typeof usage.total_tokens !== "number") {
+    return null;
+  }
+
+  return (
+    <div className={"flex items-center gap-3 " + (className ?? "")}>
+      <div className="text-sm">
+        <span className="text-muted-foreground">Total</span>{" "}
+        <span className="font-medium">{usage.total_tokens.toLocaleString()}</span>
+        <span className="text-muted-foreground"> tokens</span>
+      </div>
+    </div>
+  );
+}
+
+export default TokenCountInfo;

src/hooks/codexEvents/sessionHandlers.ts

@@ -3,6 +3,7 @@ import { toast } from "sonner";
 import { generateUniqueId } from "@/utils/genUniqueId";
 import type { ChatMessage } from "@/types/chat";
 import type { CodexEventHandler } from "./types";
+import { useEphemeralStore } from "@/stores/EphemeralStore";
 
 const handleSessionConfigured: CodexEventHandler = (event, context) => {
   if (event.msg.type !== "session_configured") {
@@ -178,11 +179,13 @@ const handleTurnAborted: CodexEventHandler = (event, context) => {
   }
 };
 
-const handleTokenCount: CodexEventHandler = (event) => {
+const handleTokenCount: CodexEventHandler = (event, context) => {
   if (event.msg.type !== "token_count") {
     return;
   }
-  console.log("token_count", event.msg);
+  const { sessionId } = context;
+  const usage = event.msg.info?.total_token_usage;
+  useEphemeralStore.getState().setSessionTokenUsage(sessionId, usage ?? undefined);
 };
 
 const handleStreamError: CodexEventHandler = (event) => {

src/stores/EphemeralStore.ts

@@ -1,14 +1,26 @@
 import { create } from 'zustand';
 
+interface TokenCountShape {
+  input_tokens?: number;
+  cached_input_tokens?: number;
+  output_tokens?: number;
+  reasoning_output_tokens?: number;
+  total_tokens: number;
+}
+
 interface EphemeralStoreState {
   // sessionId -> filePath -> { unified, updatedAt }
   sessionFileDiffs: Record<string, Record<string, { unified: string; updatedAt: number }>>;
+  // sessionId -> latest token usage
+  sessionTokenUsage: Record<string, TokenCountShape | undefined>;
   setTurnDiff: (sessionId: string, unifiedDiff: string) => void;
   clearTurnDiffs: (sessionId: string) => void;
+  setSessionTokenUsage: (sessionId: string, usage?: TokenCountShape) => void;
 }
 
 export const useEphemeralStore = create<EphemeralStoreState>((set) => ({
   sessionFileDiffs: {},
+  sessionTokenUsage: {},
   setTurnDiff: (sessionId, unifiedDiff) =>
     set((state) => {
       // Parse unified diff into per-file segments and merge (overwrite by file)
@@ -65,4 +77,8 @@ export const useEphemeralStore = create<EphemeralStoreState>((set) => ({
       delete next[sessionId];
       return { sessionFileDiffs: next };
     }),
+  setSessionTokenUsage: (sessionId, usage) =>
+    set((state) => ({
+      sessionTokenUsage: { ...state.sessionTokenUsage, [sessionId]: usage },
+    })),
 }));