From d27511dad0c2776528c1200f427147ae067dcda8 Mon Sep 17 00:00:00 2001 From: Agung Subastian Date: Thu, 14 May 2026 19:47:27 +0700 Subject: [PATCH 1/2] fix: detect incomplete Anthropic SSE stream and deduplicate system prompt MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit stream_anthropic_sse previously returned partial output silently when the connection dropped before message_stop arrived. Now tracks message_stop_received and returns an explicit error so the user knows to retry rather than seeing a silently truncated response. Also propagates JSON parse errors in parse_anthropic_sse_line via ? instead of silently returning Ok(false), consistent with the OpenAI SSE parser. Extracts the duplicated CHAT_SYSTEM_INSTRUCTIONS into a single module- level constant — previously maintained as two identical inline concat! blocks in send_openai_compatible and send_anthropic. --- src-tauri/src/services/chat_service.rs | 51 ++++++++++++-------------- 1 file changed, 24 insertions(+), 27 deletions(-) diff --git a/src-tauri/src/services/chat_service.rs b/src-tauri/src/services/chat_service.rs index 36f09bd..e9a3330 100644 --- a/src-tauri/src/services/chat_service.rs +++ b/src-tauri/src/services/chat_service.rs @@ -14,6 +14,16 @@ 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 `` 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." +); + pub async fn get_messages(db: &SqlitePool, session_id: &str) -> AppResult> { let messages = sqlx::query_as::<_, Message>( "SELECT id, session_id, role, content, created_at FROM messages \ @@ -180,22 +190,13 @@ async fn send_openai_compatible( "stream": 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 `` 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(); @@ -250,19 +251,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 `` 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 @@ -372,8 +364,9 @@ async fn stream_anthropic_sse( let mut stream = response.bytes_stream(); let mut line_buffer = String::new(); let mut output = String::new(); + let mut message_stop_received = false; - loop { + 'outer: loop { tokio::select! { _ = cancel_token.cancelled() => { return Err(AppError::Cancelled); @@ -391,7 +384,8 @@ async fn stream_anthropic_sse( } if parse_anthropic_sse_line(&line, on_token, &mut output)? { - return Ok(output); + message_stop_received = true; + break 'outer; } } } @@ -402,6 +396,12 @@ async fn stream_anthropic_sse( } } + if !message_stop_received { + return Err(AppError::Http( + "Stream ended without completion signal — connection may have been interrupted. Please retry.".to_string(), + )); + } + Ok(output) } @@ -428,10 +428,7 @@ 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(""); From aa5152ed9950ae5e1e17b721f349f7b710baef54 Mon Sep 17 00:00:00 2001 From: Agung Subastian Date: Thu, 14 May 2026 20:09:35 +0700 Subject: [PATCH 2/2] feat: add token usage counter per session Tracks prompt and completion token counts from both OpenAI-compatible and Anthropic SSE streams and accumulates them per session in the frontend store. Backend: - TokenUsage / ChatUsageEvent / UsageAccumulator structs in chat_service.rs using serde::Serialize (no serde_json::json! to avoid clippy::disallowed_methods) - stream_openai_sse: requests stream_options.include_usage=true so the final SSE chunk carries usage; parsed in parse_openai_sse_line - stream_anthropic_sse: captures input_tokens from message_start and output_tokens from message_delta events - Emits chat-usage Tauri event after each completed completion - Also fixes stream_anthropic_sse to return error on missing message_stop (same as the pending PR #26) Frontend: - TokenUsage / ChatUsageEvent types added to types/index.ts - useChatStore: sessionUsage record, addTokenUsage (cumulative per-session sum), clearSessionUsage - AppShell: listens for chat-usage, calls addTokenUsage - ChatHeader: shows total token badge next to session title; tooltip shows split prompt / completion counts; formats as 4.2k for readability --- src-tauri/src/services/chat_service.rs | 99 +++++++++++++++++++++++--- src/components/layout/AppShell.tsx | 9 ++- src/components/layout/ChatHeader.tsx | 13 ++++ src/stores/useChatStore.ts | 25 ++++++- src/types/index.ts | 11 +++ 5 files changed, 143 insertions(+), 14 deletions(-) diff --git a/src-tauri/src/services/chat_service.rs b/src-tauri/src/services/chat_service.rs index e9a3330..91dd7ec 100644 --- a/src-tauri/src/services/chat_service.rs +++ b/src-tauri/src/services/chat_service.rs @@ -24,6 +24,42 @@ const CHAT_SYSTEM_INSTRUCTIONS: &str = concat!( "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 { + 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> { let messages = sqlx::query_as::<_, Message>( "SELECT id, session_id, role, content, created_at FROM messages \ @@ -128,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( @@ -142,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(), @@ -172,7 +215,7 @@ async fn send_openai_compatible( history: Vec, on_token: &Channel, cancel_token: &CancellationToken, -) -> AppResult { +) -> AppResult<(String, Option)> { let client = reqwest::Client::new(); let endpoint = format!("{}/chat/completions", base_url.trim_end_matches('/')); @@ -188,6 +231,7 @@ async fn send_openai_compatible( "presence_penalty": 0.0, "frequency_penalty": 0.0, "stream": true, + "stream_options": { "include_usage": true }, }); let payload_with_system = if let Some(arr) = payload.get("messages").and_then(Value::as_array) { @@ -231,7 +275,7 @@ async fn send_anthropic( api_key: Option<&str>, on_token: &Channel, cancel_token: &CancellationToken, -) -> AppResult { +) -> AppResult<(String, Option)> { let client = reqwest::Client::new(); let (system_msgs, chat_msgs): (Vec<_>, Vec<_>) = @@ -281,10 +325,11 @@ async fn stream_openai_sse( response: reqwest::Response, on_token: &Channel, cancel_token: &CancellationToken, -) -> AppResult { +) -> AppResult<(String, Option)> { 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! { @@ -303,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())); } } } @@ -316,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, output: &mut String, + usage: &mut UsageAccumulator, ) -> AppResult { let trimmed = line.trim(); if trimmed.is_empty() { @@ -341,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) @@ -360,10 +416,11 @@ async fn stream_anthropic_sse( response: reqwest::Response, on_token: &Channel, cancel_token: &CancellationToken, -) -> AppResult { +) -> AppResult<(String, Option)> { 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; 'outer: loop { @@ -383,7 +440,7 @@ async fn stream_anthropic_sse( line.pop(); } - if parse_anthropic_sse_line(&line, on_token, &mut output)? { + if parse_anthropic_sse_line(&line, on_token, &mut output, &mut usage)? { message_stop_received = true; break 'outer; } @@ -402,13 +459,14 @@ async fn stream_anthropic_sse( )); } - Ok(output) + Ok((output, usage.finish())) } fn parse_anthropic_sse_line( line: &str, on_token: &Channel, output: &mut String, + usage: &mut UsageAccumulator, ) -> AppResult { let trimmed = line.trim(); if trimmed.is_empty() { @@ -433,6 +491,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") diff --git a/src/components/layout/AppShell.tsx b/src/components/layout/AppShell.tsx index adbdfdf..7eb5343 100644 --- a/src/components/layout/AppShell.tsx +++ b/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('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, diff --git a/src/components/layout/ChatHeader.tsx b/src/components/layout/ChatHeader.tsx index df8094f..5d4cbf1 100644 --- a/src/components/layout/ChatHeader.tsx +++ b/src/components/layout/ChatHeader.tsx @@ -25,9 +25,14 @@ export const ChatHeader: React.FC = ({ 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 = ({ onToggleLeftSidebar }) = )} + {currentUsage && ( + + {formatTokens(currentUsage.totalTokens)} tokens + + )} )} diff --git a/src/stores/useChatStore.ts b/src/stores/useChatStore.ts index 2256568..0d29500 100644 --- a/src/stores/useChatStore.ts +++ b/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; 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((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 }; + }), })); diff --git a/src/types/index.ts b/src/types/index.ts index 4c26385..a17cc99 100644 --- a/src/types/index.ts +++ b/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;