From be8cd7b6ebfb64d08c0b1c2f6449f19e1703de73 Mon Sep 17 00:00:00 2001
From: Jwalin Shah <jwalinshah13@gmail.com>
Date: Tue, 21 Apr 2026 14:40:08 -0700
Subject: [PATCH] feat(chat): show active token consumption and plan status in
 UI (#703) (#738)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Adds a small pill in the Conversations chat header with:
- Session token counter — cumulative input + output tokens across all
  chat turns, accumulated from chat:done events via new
  recordChatTurnUsage slice action.
- Plan usage badge — % of 5-hour or weekly budget consumed, driven by
  existing useUsageState. Color-coded sage/amber/coral by severity;
  click navigates to /settings/billing.

No new backend endpoints — session counter is client-local, plan state
reuses the existing creditsApi/billingApi fetches cached by
useUsageState.

Closes #703.

Co-authored-by: Jwalin Shah <jshah1331@gmail.com>
---
 app/src/components/chat/TokenUsagePill.tsx | 95 ++++++++++++++++++++++
 app/src/pages/Conversations.tsx            |  2 +
 app/src/providers/ChatRuntimeProvider.tsx  |  9 ++
 app/src/store/chatRuntimeSlice.ts          | 24 ++++++
 4 files changed, 130 insertions(+)
 create mode 100644 app/src/components/chat/TokenUsagePill.tsx

diff --git a/app/src/components/chat/TokenUsagePill.tsx b/app/src/components/chat/TokenUsagePill.tsx
new file mode 100644
index 000000000..50badfbeb
--- /dev/null
+++ b/app/src/components/chat/TokenUsagePill.tsx
@@ -0,0 +1,95 @@
+import { useNavigate } from 'react-router-dom';
+
+import { useUsageState } from '../../hooks/useUsageState';
+import { useAppSelector } from '../../store/hooks';
+
+function formatTokens(n: number): string {
+  if (n < 1000) return String(n);
+  if (n < 1_000_000) return `${(n / 1000).toFixed(n < 10_000 ? 1 : 0)}K`;
+  return `${(n / 1_000_000).toFixed(1)}M`;
+}
+
+interface PillSeverity {
+  bg: string;
+  text: string;
+  ring: string;
+  label: string;
+}
+
+function severityFromPct(pct: number): PillSeverity {
+  if (pct >= 0.9) {
+    return {
+      bg: 'bg-coral-50',
+      text: 'text-coral-700',
+      ring: 'ring-coral-200',
+      label: `${Math.round(pct * 100)}%`,
+    };
+  }
+  if (pct >= 0.7) {
+    return {
+      bg: 'bg-amber-50',
+      text: 'text-amber-700',
+      ring: 'ring-amber-200',
+      label: `${Math.round(pct * 100)}%`,
+    };
+  }
+  return {
+    bg: 'bg-sage-50',
+    text: 'text-sage-700',
+    ring: 'ring-sage-200',
+    label: `${Math.round(pct * 100)}%`,
+  };
+}
+
+const TokenUsagePill = () => {
+  const navigate = useNavigate();
+  const sessionTokens = useAppSelector(state => state.chatRuntime.sessionTokenUsage);
+  const { usagePct10h, usagePct7d, isAtLimit, isNearLimit, currentTier, teamUsage } =
+    useUsageState();
+
+  const totalTokens = sessionTokens.inputTokens + sessionTokens.outputTokens;
+  const showSessionCounter = totalTokens > 0;
+
+  const planPct = Math.max(usagePct10h, usagePct7d);
+  const planSeverity = severityFromPct(planPct);
+  const showPlanPill = teamUsage !== null;
+
+  const planTitle = (() => {
+    if (isAtLimit) return 'Usage limit reached — click to top up';
+    if (isNearLimit) return 'Approaching usage limit';
+    return `${currentTier.toLowerCase()} plan — click for details`;
+  })();
+
+  if (!showSessionCounter && !showPlanPill) return null;
+
+  return (
+    <div className="flex items-center gap-1.5 text-[11px] leading-none">
+      {showSessionCounter ? (
+        <span
+          className="inline-flex items-center gap-1 rounded-full bg-stone-100 px-2 py-1 font-mono text-stone-600 ring-1 ring-stone-200/60"
+          title={`Session tokens: ${sessionTokens.inputTokens.toLocaleString()} in / ${sessionTokens.outputTokens.toLocaleString()} out across ${sessionTokens.turns} turn(s)`}>
+          <svg className="h-3 w-3" fill="none" stroke="currentColor" viewBox="0 0 24 24">
+            <path
+              strokeLinecap="round"
+              strokeLinejoin="round"
+              strokeWidth={2}
+              d="M13 10V3L4 14h7v7l9-11h-7z"
+            />
+          </svg>
+          {formatTokens(totalTokens)}
+        </span>
+      ) : null}
+      {showPlanPill ? (
+        <button
+          type="button"
+          onClick={() => navigate('/settings/billing')}
+          title={planTitle}
+          className={`inline-flex items-center gap-1 rounded-full px-2 py-1 font-medium ring-1 transition-colors ${planSeverity.bg} ${planSeverity.text} ${planSeverity.ring} hover:opacity-80`}>
+          {isAtLimit ? 'Limit' : planSeverity.label}
+        </button>
+      ) : null}
+    </div>
+  );
+};
+
+export default TokenUsagePill;
diff --git a/app/src/pages/Conversations.tsx b/app/src/pages/Conversations.tsx
index 850c922d4..ef625f896 100644
--- a/app/src/pages/Conversations.tsx
+++ b/app/src/pages/Conversations.tsx
@@ -4,6 +4,7 @@ import Markdown from 'react-markdown';
 import { useNavigate } from 'react-router-dom';
 
 import { type ChatSendError, chatSendError } from '../chat/chatSendError';
+import TokenUsagePill from '../components/chat/TokenUsagePill';
 import UpsellBanner from '../components/upsell/UpsellBanner';
 import { dismissBanner, shouldShowBanner } from '../components/upsell/upsellDismissState';
 import UsageLimitModal from '../components/upsell/UsageLimitModal';
@@ -1114,6 +1115,7 @@ const Conversations = ({ variant = 'page' }: ConversationsProps = {}) => {
             <h3 className="text-sm font-medium text-stone-700 truncate flex-1">
               {threads.find(t => t.id === selectedThreadId)?.title ?? 'Select a thread'}
             </h3>
+            <TokenUsagePill />
             <button
               onClick={() => void handleCreateNewThread()}
               className="px-2.5 py-1 rounded-lg text-xs font-medium text-primary-600 hover:bg-primary-50 transition-colors"
diff --git a/app/src/providers/ChatRuntimeProvider.tsx b/app/src/providers/ChatRuntimeProvider.tsx
index 769f9fef9..52ac11cf4 100644
--- a/app/src/providers/ChatRuntimeProvider.tsx
+++ b/app/src/providers/ChatRuntimeProvider.tsx
@@ -20,6 +20,7 @@ import {
   clearStreamingAssistantForThread,
   endInferenceTurn,
   markInferenceTurnStreaming,
+  recordChatTurnUsage,
   setInferenceStatusForThread,
   setStreamingAssistantForThread,
   setToolTimelineForThread,
@@ -499,8 +500,16 @@ const ChatRuntimeProvider = ({ children }: { children: React.ReactNode }) => {
           thread: event.thread_id,
           request: event.request_id,
           segments: event.segment_total,
+          input_tokens: event.total_input_tokens,
+          output_tokens: event.total_output_tokens,
         });
 
+        dispatch(
+          recordChatTurnUsage({
+            inputTokens: event.total_input_tokens,
+            outputTokens: event.total_output_tokens,
+          })
+        );
         dispatch(clearInferenceStatusForThread({ threadId: event.thread_id }));
         dispatch(clearStreamingAssistantForThread({ threadId: event.thread_id }));
 
diff --git a/app/src/store/chatRuntimeSlice.ts b/app/src/store/chatRuntimeSlice.ts
index 312d364f3..58598f401 100644
--- a/app/src/store/chatRuntimeSlice.ts
+++ b/app/src/store/chatRuntimeSlice.ts
@@ -36,6 +36,14 @@ export interface StreamingAssistantState {
  */
 export type InferenceTurnLifecycle = 'started' | 'streaming';
 
+/** Running per-session totals accumulated from `chat:done` events (#703). */
+export interface SessionTokenUsage {
+  inputTokens: number;
+  outputTokens: number;
+  turns: number;
+  lastUpdated: number;
+}
+
 /**
  * Per-thread UI state for an in-flight agent turn (socket events while the user
  * may navigate away from Conversations). The thread slice keeps `activeThreadId`
@@ -47,6 +55,7 @@ interface ChatRuntimeState {
   streamingAssistantByThread: Record<string, StreamingAssistantState>;
   toolTimelineByThread: Record<string, ToolTimelineEntry[]>;
   inferenceTurnLifecycleByThread: Record<string, InferenceTurnLifecycle>;
+  sessionTokenUsage: SessionTokenUsage;
 }
 
 const initialState: ChatRuntimeState = {
@@ -54,6 +63,7 @@ const initialState: ChatRuntimeState = {
   streamingAssistantByThread: {},
   toolTimelineByThread: {},
   inferenceTurnLifecycleByThread: {},
+  sessionTokenUsage: { inputTokens: 0, outputTokens: 0, turns: 0, lastUpdated: 0 },
 };
 
 const chatRuntimeSlice = createSlice({
@@ -110,6 +120,18 @@ const chatRuntimeSlice = createSlice({
       state.toolTimelineByThread = {};
       state.inferenceTurnLifecycleByThread = {};
     },
+    recordChatTurnUsage: (
+      state,
+      action: PayloadAction<{ inputTokens: number; outputTokens: number }>
+    ) => {
+      state.sessionTokenUsage.inputTokens += Math.max(0, action.payload.inputTokens);
+      state.sessionTokenUsage.outputTokens += Math.max(0, action.payload.outputTokens);
+      state.sessionTokenUsage.turns += 1;
+      state.sessionTokenUsage.lastUpdated = Date.now();
+    },
+    resetSessionTokenUsage: state => {
+      state.sessionTokenUsage = { inputTokens: 0, outputTokens: 0, turns: 0, lastUpdated: 0 };
+    },
   },
 });
 
@@ -125,6 +147,8 @@ export const {
   endInferenceTurn,
   clearRuntimeForThread,
   clearAllChatRuntime,
+  recordChatTurnUsage,
+  resetSessionTokenUsage,
 } = chatRuntimeSlice.actions;
 
 export default chatRuntimeSlice.reducer;