feat: add token usage counter per session

src-tauri/src/services/chat_service.rs

@@ -14,6 +14,52 @@ use crate::{
 
 use super::{now_rfc3339, provider_service};
 
+const CHAT_SYSTEM_INSTRUCTIONS: &str = concat!(
+    "IMPORTANT: Reply using the same language as the user's latest message. If user writes Indonesian, answer in Indonesian. Never switch to another language unless the user explicitly asks you to.\n\n",
+    "INTERACTIVE PREVIEW: When the user asks for a visualization, diagram, chart, interactive demo, or any visual HTML content, output it as a fenced code block with tag `html:preview`. The app renders it as a live iframe preview with a full design system pre-loaded (CSS variables, SVG color ramp classes, pre-styled form elements, light/dark mode).\n\n",
+    "Design rules: flat (no gradients/shadows/glow), use CSS vars for colors (var(--color-text-primary), var(--color-background-secondary), etc). system-ui font, 2 weights (400/500), sentence case. Structure: style → content → script last.\n\n",
+    "SVG diagrams: use pre-loaded classes — `.t` (14px text), `.ts` (12px), `.th` (14px bold), `.box` (neutral), `.node` (clickable), `.arr` (arrow), `.leader` (dashed). Color ramps: `class=\"c-blue\"` on `<g>` wrapping shape+text — auto light/dark. Available: c-purple, c-teal, c-coral, c-blue, c-amber, c-green, c-red, c-gray, c-pink. Max 2-3 ramps per diagram.\n\n",
+    "Chart.js: wrap canvas in div with position:relative + explicit height. Load UMD from cdnjs.cloudflare.com with onload callback. Disable default legend, build custom HTML legend with 10px colored squares.\n\n",
+    "Interactive: form elements pre-styled. Use sendPrompt(text) for drill-down. CDN: cdnjs.cloudflare.com, cdn.jsdelivr.net, unpkg.com, esm.sh only.\n\n",
+    "Always output COMPLETE standalone HTML (DOCTYPE, html, head, body). No titles/prose inside widget — explanations go in your response text."
+);
+
+#[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 +164,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 +178,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 +215,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,24 +231,16 @@ async fn send_openai_compatible(
         "presence_penalty": 0.0,
         "frequency_penalty": 0.0,
         "stream": true,
+        "stream_options": { "include_usage": true },
     });
 
-    let system_instructions = concat!(
-        "IMPORTANT: Reply using the same language as the user's latest message. If user writes Indonesian, answer in Indonesian. Never switch to another language unless the user explicitly asks you to.\n\n",
-        "INTERACTIVE PREVIEW: When the user asks for a visualization, diagram, chart, interactive demo, or any visual HTML content, output it as a fenced code block with tag `html:preview`. The app renders it as a live iframe preview with a full design system pre-loaded (CSS variables, SVG color ramp classes, pre-styled form elements, light/dark mode).\n\n",
-        "Design rules: flat (no gradients/shadows/glow), use CSS vars for colors (var(--color-text-primary), var(--color-background-secondary), etc). system-ui font, 2 weights (400/500), sentence case. Structure: style → content → script last.\n\n",
-        "SVG diagrams: use pre-loaded classes — `.t` (14px text), `.ts` (12px), `.th` (14px bold), `.box` (neutral), `.node` (clickable), `.arr` (arrow), `.leader` (dashed). Color ramps: `class=\"c-blue\"` on `<g>` wrapping shape+text — auto light/dark. Available: c-purple, c-teal, c-coral, c-blue, c-amber, c-green, c-red, c-gray, c-pink. Max 2-3 ramps per diagram.\n\n",
-        "Chart.js: wrap canvas in div with position:relative + explicit height. Load UMD from cdnjs.cloudflare.com with onload callback. Disable default legend, build custom HTML legend with 10px colored squares.\n\n",
-        "Interactive: form elements pre-styled. Use sendPrompt(text) for drill-down. CDN: cdnjs.cloudflare.com, cdn.jsdelivr.net, unpkg.com, esm.sh only.\n\n",
-        "Always output COMPLETE standalone HTML (DOCTYPE, html, head, body). No titles/prose inside widget — explanations go in your response text."
-    );
     let payload_with_system = if let Some(arr) = payload.get("messages").and_then(Value::as_array) {
         let mut updated = arr.clone();
         updated.insert(
             0,
             serde_json::json!({
                 "role": "system",
-                "content": system_instructions,
+                "content": CHAT_SYSTEM_INSTRUCTIONS,
             }),
         );
         let mut p = payload.clone();
@@ -230,7 +275,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<_>) =
@@ -250,19 +295,10 @@ async fn send_anthropic(
         "stream": true,
     });
 
-    let system_instructions_anthropic = concat!(
-        "IMPORTANT: Reply using the same language as the user's latest message. If user writes Indonesian, answer in Indonesian. Never switch to another language unless the user explicitly asks you to.\n\n",
-        "INTERACTIVE PREVIEW: When the user asks for a visualization, diagram, chart, interactive demo, or any visual HTML content, output it as a fenced code block with tag `html:preview`. The app renders it as a live iframe preview with a full design system pre-loaded (CSS variables, SVG color ramp classes, pre-styled form elements, light/dark mode).\n\n",
-        "Design rules: flat (no gradients/shadows/glow), use CSS vars for colors (var(--color-text-primary), var(--color-background-secondary), etc). system-ui font, 2 weights (400/500), sentence case. Structure: style → content → script last.\n\n",
-        "SVG diagrams: use pre-loaded classes — `.t` (14px text), `.ts` (12px), `.th` (14px bold), `.box` (neutral), `.node` (clickable), `.arr` (arrow), `.leader` (dashed). Color ramps: `class=\"c-blue\"` on `<g>` wrapping shape+text — auto light/dark. Available: c-purple, c-teal, c-coral, c-blue, c-amber, c-green, c-red, c-gray, c-pink. Max 2-3 ramps per diagram.\n\n",
-        "Chart.js: wrap canvas in div with position:relative + explicit height. Load UMD from cdnjs.cloudflare.com with onload callback. Disable default legend, build custom HTML legend with 10px colored squares.\n\n",
-        "Interactive: form elements pre-styled. Use sendPrompt(text) for drill-down. CDN: cdnjs.cloudflare.com, cdn.jsdelivr.net, unpkg.com, esm.sh only.\n\n",
-        "Always output COMPLETE standalone HTML (DOCTYPE, html, head, body). No titles/prose inside widget — explanations go in your response text."
-    );
     if let Some(sys) = system_msgs.first() {
-        payload["system"] = serde_json::json!(format!("{}\n\n{}", sys.content, system_instructions_anthropic));
+        payload["system"] = serde_json::json!(format!("{}\n\n{}", sys.content, CHAT_SYSTEM_INSTRUCTIONS));
     } else {
-        payload["system"] = serde_json::json!(system_instructions_anthropic);
+        payload["system"] = serde_json::json!(CHAT_SYSTEM_INSTRUCTIONS);
     }
 
     let mut request = client
@@ -289,10 +325,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 +348,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 +361,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 +387,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 +416,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 +440,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 +453,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() {
@@ -428,14 +486,30 @@ fn parse_anthropic_sse_line(
     };
     let payload = payload.trim();
 
-    let value: Value = match serde_json::from_str(payload) {
-        Ok(v) => v,
-        Err(_) => return Ok(false),
-    };
+    let value: Value = serde_json::from_str(payload)?;
 
     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>