🤖 Refine IPC: Move token consumer calculation to frontend

Inverts ownership for better separation of concerns:
- Backend: Just tokenization (single tokens:countBulk endpoint)
- Frontend: Display logic (aggregation, percentages, sorting)

Benefits:
- Simpler API (one endpoint vs separate stats.calculate)
- Pure functions = easier testing (no IPC mocking needed)
- Frontend can batch all tokenization in one call
- Clear responsibilities: backend tokenizes, frontend calculates

Changes:
- Added tokens:countBulk IPC endpoint
- Created consumerCalculator.ts with pure calculation functions
- Updated TokenConsumerBreakdown to use new endpoint
- Added comprehensive unit tests (12 test cases)
- Removed unused styled components from CostsTab

Stats calculation now happens entirely in the frontend using raw
tokenization results from the backend. The old stats.calculate
endpoint remains for backward compatibility but can be deprecated.

Author: ammar-agent

src/components/ChatMetaSidebar/CostsTab.tsx

@@ -20,15 +20,6 @@ const Section = styled.div`
   margin-bottom: 24px;
 `;
 
-const SectionTitle = styled.h3<{ dimmed?: boolean }>`
-  color: ${(props) => (props.dimmed ? "#999999" : "#cccccc")};
-  font-size: 14px;
-  font-weight: 600;
-  margin: 0 0 12px 0;
-  text-transform: uppercase;
-  letter-spacing: 0.5px;
-`;
-
 const ConsumerList = styled.div`
   display: flex;
   flex-direction: column;
@@ -116,11 +107,6 @@ const CachedSegment = styled.div<SegmentProps>`
   transition: width 0.3s ease;
 `;
 
-const LoadingState = styled.div`
-  color: #888888;
-  font-style: italic;
-`;
-
 const EmptyState = styled.div`
   color: #888888;
   text-align: center;
@@ -538,7 +524,7 @@ export const CostsTab: React.FC<CostsTabProps> = ({ workspaceId }) => {
         </Section>
       )}
 
-      <TokenConsumerBreakdown workspaceId={workspaceId} model={model} />
+      <TokenConsumerBreakdown messages={messages} model={model} />
     </Container>
   );
 };

src/components/ChatMetaSidebar/TokenConsumerBreakdown.tsx

@@ -1,6 +1,8 @@
 import React, { useState, useEffect } from "react";
 import styled from "@emotion/styled";
 import type { ChatStats } from "@/types/chatStats";
+import type { CmuxMessage } from "@/types/message";
+import { prepareTokenization, calculateConsumers } from "@/utils/tokens/consumerCalculator";
 
 const Section = styled.div`
   margin-bottom: 24px;
@@ -102,12 +104,12 @@ const formatTokens = (tokens: number): string => {
 };
 
 interface TokenConsumerBreakdownProps {
-  workspaceId: string;
+  messages: CmuxMessage[];
   model: string;
 }
 
 export const TokenConsumerBreakdown: React.FC<TokenConsumerBreakdownProps> = ({
-  workspaceId,
+  messages,
   model,
 }) => {
   const [stats, setStats] = useState<ChatStats | null>(null);
@@ -116,27 +118,61 @@ export const TokenConsumerBreakdown: React.FC<TokenConsumerBreakdownProps> = ({
   useEffect(() => {
     let cancelled = false;
 
-    setIsCalculating(true);
-    window.api.stats
-      .calculate(workspaceId, model)
-      .then((result) => {
-        if (!cancelled && result) {
-          setStats(result);
+    async function calculate() {
+      setIsCalculating(true);
+
+      try {
+        // Prepare all text for tokenization (pure function)
+        const { texts, consumerMap, toolDefinitions } = prepareTokenization(messages, model);
+
+        // Combine message texts + tool definition strings for bulk tokenization
+        const allTexts = [...texts, ...Array.from(toolDefinitions.values())];
+
+        // Batch tokenize everything in one IPC call
+        const tokenCounts = await window.api.tokens.countBulk(model, allTexts);
+
+        if (cancelled || !tokenCounts) {
+          return; // Tokenizer not loaded or component unmounted
         }
-      })
-      .catch((error) => {
+
+        // Split results back into message tokens and tool definition tokens
+        const messageTokens = tokenCounts.slice(0, texts.length);
+        const toolDefCounts = new Map<string, number>();
+        let defIndex = texts.length;
+        for (const [toolName] of toolDefinitions) {
+          toolDefCounts.set(toolName, tokenCounts[defIndex]);
+          defIndex++;
+        }
+
+        // Calculate consumers (pure function)
+        const consumers = calculateConsumers(messageTokens, consumerMap, toolDefCounts);
+        const totalTokens = consumers.reduce((sum, c) => sum + c.tokens, 0);
+
+        // Derive tokenizer name from model
+        const tokenizerName = model.startsWith("anthropic:") ? "claude" : "o200k_base";
+
+        setStats({
+          consumers,
+          totalTokens,
+          model,
+          tokenizerName,
+          usageHistory: [], // Not used in this component
+        });
+      } catch (error) {
         console.error(`[TokenConsumerBreakdown] Failed to calculate stats:`, error);
-      })
-      .finally(() => {
+      } finally {
         if (!cancelled) {
           setIsCalculating(false);
         }
-      });
+      }
+    }
+
+    void calculate();
 
     return () => {
       cancelled = true;
     };
-  }, [workspaceId, model]);
+  }, [messages, model]);
 
   if (isCalculating) {
     return (

src/constants/ipc-constants.ts

@@ -41,6 +41,9 @@ export const IPC_CHANNELS = {
   // Stats channels
   STATS_CALCULATE: "stats:calculate",
 
+  // Token channels
+  TOKENS_COUNT_BULK: "tokens:countBulk",
+
   // Dynamic channel prefixes
   WORKSPACE_CHAT_PREFIX: "workspace:chat:",
   WORKSPACE_METADATA: "workspace:metadata",

src/preload.ts

@@ -114,6 +114,10 @@ const api: IPCApi = {
     calculate: (workspaceId: string, model: string) =>
       ipcRenderer.invoke(IPC_CHANNELS.STATS_CALCULATE, workspaceId, model),
   },
+  tokens: {
+    countBulk: (model: string, texts: string[]) =>
+      ipcRenderer.invoke(IPC_CHANNELS.TOKENS_COUNT_BULK, model, texts),
+  },
 };
 
 // Expose the API along with platform/versions

src/services/ipcMain.ts

@@ -141,6 +141,7 @@ export class IpcMain {
     this.registerDialogHandlers(ipcMain);
     this.registerWindowHandlers(ipcMain);
     this.registerStatsHandlers(ipcMain);
+    this.registerTokenHandlers(ipcMain);
     this.registerWorkspaceHandlers(ipcMain);
     this.registerProviderHandlers(ipcMain);
     this.registerProjectHandlers(ipcMain);
@@ -199,6 +200,25 @@ export class IpcMain {
     );
   }
 
+  private registerTokenHandlers(ipcMain: ElectronIpcMain): void {
+    ipcMain.handle(
+      IPC_CHANNELS.TOKENS_COUNT_BULK,
+      async (_event, model: string, texts: string[]) => {
+        try {
+          // Dynamic import to lazy-load tokenizer (prevents ~3-4s startup freeze)
+          /* eslint-disable no-restricted-syntax */
+          const { getTokenizerForModel } = await import("@/utils/main/tokenizer");
+          /* eslint-enable no-restricted-syntax */
+          const tokenizer = getTokenizerForModel(model);
+          return texts.map((text) => tokenizer.countTokens(text));
+        } catch (error) {
+          log.error(`Failed to count tokens for model ${model}:`, error);
+          return null; // Tokenizer not loaded or error occurred
+        }
+      }
+    );
+  }
+
   private registerWorkspaceHandlers(ipcMain: ElectronIpcMain): void {
     ipcMain.handle(
       IPC_CHANNELS.WORKSPACE_CREATE,

src/types/ipc.ts

@@ -234,4 +234,7 @@ export interface IPCApi {
   stats: {
     calculate(workspaceId: string, model: string): Promise<ChatStats | null>;
   };
+  tokens: {
+    countBulk(model: string, texts: string[]): Promise<number[] | null>;
+  };
 }