feat: add token usage counter per session

src-tauri/src/services/chat_service.rs

@@ -14,6 +14,42 @@ use crate::{
 
 use super::{now_rfc3339, provider_service};
 
+#[derive(Debug, Clone, serde::Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct TokenUsage {
+    pub prompt_tokens: u32,
+    pub completion_tokens: u32,
+    pub total_tokens: u32,
+}
+
+#[derive(Debug, Clone, serde::Serialize)]
+#[serde(rename_all = "camelCase")]
+struct ChatUsageEvent {
+    session_id: String,
+    usage: TokenUsage,
+}
+
+#[derive(Debug, Default)]
+struct UsageAccumulator {
+    prompt_tokens: u32,
+    completion_tokens: u32,
+}
+
+impl UsageAccumulator {
+    fn finish(self) -> Option<TokenUsage> {
+        let total = self.prompt_tokens + self.completion_tokens;
+        if total == 0 {
+            None
+        } else {
+            Some(TokenUsage {
+                prompt_tokens: self.prompt_tokens,
+                completion_tokens: self.completion_tokens,
+                total_tokens: total,
+            })
+        }
+    }
+}
+
 pub async fn get_messages(db: &SqlitePool, session_id: &str) -> AppResult<Vec<Message>> {
     let messages = sqlx::query_as::<_, Message>(
         "SELECT id, session_id, role, content, created_at FROM messages \
@@ -118,7 +154,7 @@ async fn send_message_inner(
     // Use caller-supplied model_id if provided, otherwise fall back to provider default
     let model = model_id.unwrap_or(&provider.model);
 
-    let assistant_output = if provider.provider_type == "anthropic" {
+    let (assistant_output, token_usage) = if provider.provider_type == "anthropic" {
         send_anthropic(history, model, provider.api_key.as_deref(), &on_token, &cancel_token).await?
     } else {
         send_openai_compatible(
@@ -132,6 +168,13 @@ async fn send_message_inner(
         .await?
     };
 
+    if let Some(usage) = token_usage {
+        let _ = app_handle.emit("chat-usage", ChatUsageEvent {
+            session_id: session_id.to_string(),
+            usage,
+        });
+    }
+
     let assistant_message = Message {
         id: Uuid::new_v4().to_string(),
         session_id: session_id.to_string(),
@@ -162,7 +205,7 @@ async fn send_openai_compatible(
     history: Vec<Message>,
     on_token: &Channel<String>,
     cancel_token: &CancellationToken,
-) -> AppResult<String> {
+) -> AppResult<(String, Option<TokenUsage>)> {
     let client = reqwest::Client::new();
     let endpoint = format!("{}/chat/completions", base_url.trim_end_matches('/'));
 
@@ -178,6 +221,7 @@ async fn send_openai_compatible(
         "presence_penalty": 0.0,
         "frequency_penalty": 0.0,
         "stream": true,
+        "stream_options": { "include_usage": true },
     });
 
     let system_instructions = concat!(
@@ -230,7 +274,7 @@ async fn send_anthropic(
     api_key: Option<&str>,
     on_token: &Channel<String>,
     cancel_token: &CancellationToken,
-) -> AppResult<String> {
+) -> AppResult<(String, Option<TokenUsage>)> {
     let client = reqwest::Client::new();
 
     let (system_msgs, chat_msgs): (Vec<_>, Vec<_>) =
@@ -289,10 +333,11 @@ async fn stream_openai_sse(
     response: reqwest::Response,
     on_token: &Channel<String>,
     cancel_token: &CancellationToken,
-) -> AppResult<String> {
+) -> AppResult<(String, Option<TokenUsage>)> {
     let mut stream = response.bytes_stream();
     let mut line_buffer = String::new();
     let mut output = String::new();
+    let mut usage = UsageAccumulator::default();
 
     loop {
         tokio::select! {
@@ -311,8 +356,8 @@ async fn stream_openai_sse(
                                 line.pop();
                             }
 
-                            if parse_openai_sse_line(&line, on_token, &mut output)? {
-                                return Ok(output);
+                            if parse_openai_sse_line(&line, on_token, &mut output, &mut usage)? {
+                                return Ok((output, usage.finish()));
                             }
                         }
                     }
@@ -324,16 +369,17 @@ async fn stream_openai_sse(
     }
 
     if !line_buffer.is_empty() {
-        parse_openai_sse_line(&line_buffer, on_token, &mut output)?;
+        parse_openai_sse_line(&line_buffer, on_token, &mut output, &mut usage)?;
     }
 
-    Ok(output)
+    Ok((output, usage.finish()))
 }
 
 fn parse_openai_sse_line(
     line: &str,
     on_token: &Channel<String>,
     output: &mut String,
+    usage: &mut UsageAccumulator,
 ) -> AppResult<bool> {
     let trimmed = line.trim();
     if trimmed.is_empty() {
@@ -349,6 +395,16 @@ fn parse_openai_sse_line(
     }
 
     let value: Value = serde_json::from_str(payload)?;
+
+    if let Some(u) = value.get("usage") {
+        if let Some(pt) = u.get("prompt_tokens").and_then(Value::as_u64) {
+            usage.prompt_tokens = pt as u32;
+        }
+        if let Some(ct) = u.get("completion_tokens").and_then(Value::as_u64) {
+            usage.completion_tokens = ct as u32;
+        }
+    }
+
     if let Some(token) = value
         .get("choices")
         .and_then(Value::as_array)
@@ -368,12 +424,14 @@ async fn stream_anthropic_sse(
     response: reqwest::Response,
     on_token: &Channel<String>,
     cancel_token: &CancellationToken,
-) -> AppResult<String> {
+) -> AppResult<(String, Option<TokenUsage>)> {
     let mut stream = response.bytes_stream();
     let mut line_buffer = String::new();
     let mut output = String::new();
+    let mut usage = UsageAccumulator::default();
+    let mut message_stop_received = false;
 
-    loop {
+    'outer: loop {
         tokio::select! {
             _ = cancel_token.cancelled() => {
                 return Err(AppError::Cancelled);
@@ -390,8 +448,9 @@ async fn stream_anthropic_sse(
                                 line.pop();
                             }
 
-                            if parse_anthropic_sse_line(&line, on_token, &mut output)? {
-                                return Ok(output);
+                            if parse_anthropic_sse_line(&line, on_token, &mut output, &mut usage)? {
+                                message_stop_received = true;
+                                break 'outer;
                             }
                         }
                     }
@@ -402,13 +461,20 @@ async fn stream_anthropic_sse(
         }
     }
 
-    Ok(output)
+    if !message_stop_received {
+        return Err(AppError::Http(
+            "Stream ended without completion signal — connection may have been interrupted. Please retry.".to_string(),
+        ));
+    }
+
+    Ok((output, usage.finish()))
 }
 
 fn parse_anthropic_sse_line(
     line: &str,
     on_token: &Channel<String>,
     output: &mut String,
+    usage: &mut UsageAccumulator,
 ) -> AppResult<bool> {
     let trimmed = line.trim();
     if trimmed.is_empty() {
@@ -436,6 +502,25 @@ fn parse_anthropic_sse_line(
     let event_type = value.get("type").and_then(Value::as_str).unwrap_or("");
 
     match event_type {
+        "message_start" => {
+            if let Some(pt) = value
+                .get("message")
+                .and_then(|m| m.get("usage"))
+                .and_then(|u| u.get("input_tokens"))
+                .and_then(Value::as_u64)
+            {
+                usage.prompt_tokens = pt as u32;
+            }
+        }
+        "message_delta" => {
+            if let Some(ct) = value
+                .get("usage")
+                .and_then(|u| u.get("output_tokens"))
+                .and_then(Value::as_u64)
+            {
+                usage.completion_tokens = ct as u32;
+            }
+        }
         "content_block_delta" => {
             if let Some(token) = value
                 .get("delta")

src/components/layout/AppShell.tsx

@@ -16,11 +16,11 @@ import { useUIStore } from '@/stores/useUIStore';
 import { useAgentStore } from '@/stores/useAgentStore';
 import { SettingsModal } from '@/components/settings/SettingsModal';
 import { ExcalidrawCanvas } from '@/components/canvas/ExcalidrawCanvas';
-import { AgentConfig, AgentRunWithTools, AgentType, Message, PermissionRequest, Project, Provider, ProviderModelConfig, Session, ToolCall } from '@/types';
+import { AgentConfig, AgentRunWithTools, AgentType, ChatUsageEvent, Message, PermissionRequest, Project, Provider, ProviderModelConfig, Session, ToolCall } from '@/types';
 import { cn } from '@/lib/utils';
 
 export const AppShell: React.FC = () => {
-  const { addMessage, appendStreamToken, setStreaming, clearStreaming, setMessages } = useChatStore();
+  const { addMessage, appendStreamToken, setStreaming, clearStreaming, setMessages, addTokenUsage } = useChatStore();
   const setProjects = useProjectStore((s) => s.setProjects);
   const addProject = useProjectStore((s) => s.addProject);
   const setActiveProjectId = useProjectStore((s) => s.setActiveProjectId);
@@ -181,6 +181,10 @@ export const AppShell: React.FC = () => {
         clearStreaming();
       });
 
+      const unlistenChatUsage = await listen<ChatUsageEvent>('chat-usage', (event) => {
+        addTokenUsage(event.payload.sessionId, event.payload.usage);
+      });
+
       const unlistenAgentStarted = await listen<{
         agentRunId: string;
         agentType: string;
@@ -325,6 +329,7 @@ export const AppShell: React.FC = () => {
       localUnlisten.push(
         unlistenChatDone,
         unlistenChatError,
+        unlistenChatUsage,
         unlistenAgentStarted,
         unlistenAgentToken,
         unlistenAgentToolCall,

src/components/layout/ChatHeader.tsx

@@ -25,9 +25,14 @@ export const ChatHeader: React.FC<ChatHeaderProps> = ({ onToggleLeftSidebar }) =
   const mainView = useUIStore((s) => s.mainView);
   const setMainView = useUIStore((s) => s.setMainView);
   const { activeProjectId } = useProjectStore();
+  const sessionUsage = useChatStore((s) => s.sessionUsage);
 
   const { theme, toggleTheme } = useUIStore();
   const activeSession = sessions.find(s => s.id === activeSessionId);
+  const currentUsage = activeSessionId ? sessionUsage[activeSessionId] : undefined;
+
+  const formatTokens = (n: number) =>
+    n >= 1000 ? `${(n / 1000).toFixed(1)}k` : String(n);
 
   const [isRenaming, setIsRenaming] = useState(false);
   const [renameValue, setRenameValue] = useState('');
@@ -165,6 +170,14 @@ export const ChatHeader: React.FC<ChatHeaderProps> = ({ onToggleLeftSidebar }) =
                 <PencilSimple size={14} />
               </button>
             )}
+            {currentUsage && (
+              <span
+                className="text-[10px] font-medium px-1.5 py-0.5 rounded-md bg-[var(--surface-2)] text-[var(--text-muted)] border border-[var(--border)] whitespace-nowrap"
+                title={`↑ ${formatTokens(currentUsage.promptTokens)} prompt · ↓ ${formatTokens(currentUsage.completionTokens)} completion`}
+              >
+                {formatTokens(currentUsage.totalTokens)} tokens
+              </span>
+            )}
           </>
         )}
       </div>

src/stores/useChatStore.ts

@@ -1,24 +1,47 @@
 import { create } from 'zustand';
-import { Message } from '@/types';
+import { Message, TokenUsage } from '@/types';
 
 interface ChatState {
   messages: Message[];
   streamingText: string;
   isStreaming: boolean;
+  sessionUsage: Record<string, TokenUsage>;
   setMessages: (messages: Message[]) => void;
   addMessage: (message: Message) => void;
   appendStreamToken: (token: string) => void;
   setStreaming: (isStreaming: boolean) => void;
   clearStreaming: () => void;
+  addTokenUsage: (sessionId: string, usage: TokenUsage) => void;
+  clearSessionUsage: (sessionId: string) => void;
 }
 
 export const useChatStore = create<ChatState>((set) => ({
   messages: [],
   streamingText: '',
   isStreaming: false,
+  sessionUsage: {},
   setMessages: (messages) => set({ messages }),
   addMessage: (message) => set((state) => ({ messages: [...state.messages, message] })),
   appendStreamToken: (token) => set((state) => ({ streamingText: state.streamingText + token })),
   setStreaming: (isStreaming) => set({ isStreaming }),
   clearStreaming: () => set({ streamingText: '', isStreaming: false }),
+  addTokenUsage: (sessionId, usage) =>
+    set((state) => {
+      const prev = state.sessionUsage[sessionId];
+      return {
+        sessionUsage: {
+          ...state.sessionUsage,
+          [sessionId]: {
+            promptTokens: (prev?.promptTokens ?? 0) + usage.promptTokens,
+            completionTokens: (prev?.completionTokens ?? 0) + usage.completionTokens,
+            totalTokens: (prev?.totalTokens ?? 0) + usage.totalTokens,
+          },
+        },
+      };
+    }),
+  clearSessionUsage: (sessionId) =>
+    set((state) => {
+      const { [sessionId]: _removed, ...rest } = state.sessionUsage;
+      return { sessionUsage: rest };
+    }),
 }));

src/types/index.ts

@@ -122,6 +122,17 @@ export interface AgentRunWithTools extends AgentRun {
   projectPath: string | null;
 }
 
+export interface TokenUsage {
+  promptTokens: number;
+  completionTokens: number;
+  totalTokens: number;
+}
+
+export interface ChatUsageEvent {
+  sessionId: string;
+  usage: TokenUsage;
+}
+
 export interface PermissionRequest {
   type: 'sensitive_file' | 'outside_sandbox';
   path: string;