@@ -4,7 +4,7 @@ import type { ChatMessage, ImageAttachment } from '@/components/chat';
 import type { ExtensionMessage, Model, ReactElementContext, ThinkingMode } from '@/types/protocol';
 
 import { useTauri } from '@/hooks/use-tauri';
-import { conversationAddMessage } from '@/lib/backend';
+import { conversationAddMessage, conversationLoad } from '@/lib/backend';
 import { computeSimpleDiff, getLanguageFromPath } from '@/lib/diff-utils';
 import { useFileStore } from '@/stores/file-store';
 import { useFileViewerStore } from '@/stores/file-viewer-store';
@@ -111,6 +111,7 @@ export function useChatMessages(options: UseChatMessagesOptions = {}): UseChatMe
     clearPermissions,
     addUsage,
     switchSession,
+    restoreSessionUsage,
   } = useToolStore();
   const { queueMessage: storeQueueMessage } = useQueuedMessageStore();
 
@@ -139,6 +140,81 @@ export function useChatMessages(options: UseChatMessagesOptions = {}): UseChatMe
     }
   }, [sessionId]);
 
+  // Sync tool store's currentSessionId when React sessionId changes
+  // This is critical for usage tracking - without this, usage accumulated before the first
+  // switchSession call (e.g., from localStorage restore or system:init) won't be cached
+  // when switching conversations, causing token counts to reset to 0.
+  useEffect(() => {
+    if (sessionId) {
+      const toolState = useToolStore.getState();
+      if (toolState.currentSessionId !== sessionId) {
+        switchSession(sessionId);
+      }
+    }
+  }, [sessionId, switchSession]);
+
+  // Restore usage from backend on initial mount (when sessionId comes from localStorage)
+  // This ensures token counts persist across window reloads
+  const hasRestoredUsage = useRef(false);
+  const initialSessionId = useRef(sessionId); // Capture initial sessionId
+  useEffect(() => {
+    const sid = initialSessionId.current;
+    if (!sid || hasRestoredUsage.current) return;
+
+    // Only run once on mount with initial sessionId
+    hasRestoredUsage.current = true;
+
+    // Load conversation from backend to get persisted usage data
+    void (async (): Promise<void> => {
+      try {
+        const conversation = await conversationLoad(sid);
+        if (!conversation?.messages) return;
+
+        // Calculate cumulative usage from all persisted messages
+        const processedMessageIds: string[] = [];
+        const cumulativeUsage = conversation.messages.reduce(
+          (acc, m) => {
+            if (m.usage) {
+              processedMessageIds.push(m.id);
+              return {
+                inputTokens: acc.inputTokens + m.usage.inputTokens,
+                outputTokens: acc.outputTokens + m.usage.outputTokens,
+                cacheReadInputTokens:
+                  acc.cacheReadInputTokens + (m.usage.cacheReadInputTokens ?? 0),
+                cacheCreationInputTokens:
+                  acc.cacheCreationInputTokens + (m.usage.cacheCreationInputTokens ?? 0),
+                totalCostUsd: acc.totalCostUsd + (m.usage.totalCostUsd ?? 0),
+              };
+            }
+            return acc;
+          },
+          {
+            inputTokens: 0,
+            outputTokens: 0,
+            cacheReadInputTokens: 0,
+            cacheCreationInputTokens: 0,
+            totalCostUsd: 0,
+          }
+        );
+
+        // Restore usage if we found any
+        if (processedMessageIds.length > 0) {
+          // Use getState() to avoid stale closures
+          const { restoreSessionUsage: restore, switchSession: sw } = useToolStore.getState();
+          restore(sid, cumulativeUsage, processedMessageIds);
+          // Also update current session if it matches
+          const toolState = useToolStore.getState();
+          if (toolState.currentSessionId === sid) {
+            // Re-trigger switch to apply the restored usage
+            sw(sid);
+          }
+        }
+      } catch {
+        // Ignore errors - conversation may not exist yet
+      }
+    })();
+  }, []); // Empty deps - only run on mount, uses refs for values
+
   // Track if we have pending animations - only run interval when needed
   const hasAnimatingMessage = messages.some((m) => m.displayedContent.length < m.content.length);
 
@@ -267,13 +343,31 @@ export function useChatMessages(options: UseChatMessagesOptions = {}): UseChatMe
                   : {}),
               };
 
-              // Persist assistant message to backend
+              // Persist assistant message to backend (including usage for token tracking)
+              // Build usage object conditionally to satisfy exactOptionalPropertyTypes
+              const usageDto = message.usage
+                ? {
+                    inputTokens: message.usage.input_tokens,
+                    outputTokens: message.usage.output_tokens,
+                    ...(message.usage.cache_read_input_tokens !== undefined
+                      ? { cacheReadInputTokens: message.usage.cache_read_input_tokens }
+                      : {}),
+                    ...(message.usage.cache_creation_input_tokens !== undefined
+                      ? { cacheCreationInputTokens: message.usage.cache_creation_input_tokens }
+                      : {}),
+                    ...(message.total_cost_usd !== undefined
+                      ? { totalCostUsd: message.total_cost_usd }
+                      : {}),
+                  }
+                : undefined;
+
               void conversationAddMessage(message.session_id, {
                 id: completedMsg.id,
                 role: 'assistant',
                 content: completedMsg.content,
                 ...(completedMsg.thinking ? { thinking: completedMsg.thinking } : {}),
                 createdAt: Date.now(),
+                ...(usageDto ? { usage: usageDto } : {}),
               });
 
               return [...prev.slice(0, -1), completedMsg];
@@ -376,14 +470,19 @@ export function useChatMessages(options: UseChatMessagesOptions = {}): UseChatMe
         }
 
         case 'conversation:loaded': {
-          // Prepare new messages
+          // Prepare new messages (filter out system messages as they're not displayed)
           const cachedMessages = messagesCache.current.get(message.session_id);
-          const backendMessages = message.messages.map((m) => ({
-            id: m.id,
-            role: m.role,
-            content: m.content,
-            displayedContent: m.content,
-          }));
+          const backendMessages = message.messages
+            .filter(
+              (m): m is typeof m & { role: 'user' | 'assistant' } =>
+                m.role === 'user' || m.role === 'assistant'
+            )
+            .map((m) => ({
+              id: m.id,
+              role: m.role,
+              content: m.content,
+              displayedContent: m.content,
+            }));
 
           // Determine which messages to use (prefer cache if it has more messages)
           let newMessages: typeof backendMessages;
@@ -395,6 +494,39 @@ export function useChatMessages(options: UseChatMessagesOptions = {}): UseChatMe
             newMessages = backendMessages;
           }
 
+          // Calculate cumulative usage from persisted messages (for token tracking persistence)
+          const processedMessageIds: string[] = [];
+          const cumulativeUsage = message.messages.reduce(
+            (acc, m) => {
+              if (m.usage) {
+                processedMessageIds.push(m.id);
+                return {
+                  inputTokens: acc.inputTokens + m.usage.inputTokens,
+                  outputTokens: acc.outputTokens + m.usage.outputTokens,
+                  cacheReadInputTokens:
+                    acc.cacheReadInputTokens + (m.usage.cacheReadInputTokens ?? 0),
+                  cacheCreationInputTokens:
+                    acc.cacheCreationInputTokens + (m.usage.cacheCreationInputTokens ?? 0),
+                  totalCostUsd: acc.totalCostUsd + (m.usage.totalCostUsd ?? 0),
+                };
+              }
+              return acc;
+            },
+            {
+              inputTokens: 0,
+              outputTokens: 0,
+              cacheReadInputTokens: 0,
+              cacheCreationInputTokens: 0,
+              totalCostUsd: 0,
+            }
+          );
+
+          // Restore usage from persisted data BEFORE switching session
+          // This ensures the session cache is pre-populated with correct usage
+          if (processedMessageIds.length > 0) {
+            restoreSessionUsage(message.session_id, cumulativeUsage, processedMessageIds);
+          }
+
           // Set all state atomically - React 18 batches these updates
           setMessages(newMessages);
           setSessionId(message.session_id);
@@ -607,6 +739,7 @@ export function useChatMessages(options: UseChatMessagesOptions = {}): UseChatMe
       addPermissionRequest,
       addUsage,
       switchSession,
+      restoreSessionUsage,
       onSessionCreated,
     ]
   );
 
@@ -609,7 +609,11 @@ async function handleTauriMessage(message: WebviewMessage): Promise<void> {
               id: m.id,
               role: m.role,
               content: m.content,
-              timestamp: m.createdAt, // Schema expects 'timestamp', not 'created_at'
+              createdAt: m.createdAt,
+              // Include all optional fields for backwards compatibility and usage tracking
+              ...(m.thinking ? { thinking: m.thinking } : {}),
+              ...(m.toolUses && m.toolUses.length > 0 ? { toolUses: m.toolUses } : {}),
+              ...(m.usage ? { usage: m.usage } : {}),
             })),
           },
           '*'
 
@@ -815,6 +815,7 @@ export interface ConversationMessageDto {
   thinking?: string;
   createdAt: number;
   toolUses?: ToolUseDto[];
+  usage?: TokenUsageDto;
 }
 
 export interface ToolUseDto {
@@ -825,6 +826,14 @@ export interface ToolUseDto {
   success: boolean;
 }
 
+export interface TokenUsageDto {
+  inputTokens: number;
+  outputTokens: number;
+  cacheReadInputTokens?: number;
+  cacheCreationInputTokens?: number;
+  totalCostUsd?: number;
+}
+
 export interface ConversationDto {
   sessionId: string;
   title: string;
 
@@ -118,6 +118,7 @@ export interface ToolState {
   ) => void;
   resetUsage: () => void;
   switchSession: (newSessionId: string) => void;
+  restoreSessionUsage: (sessionId: string, usage: UsageData, processedIds?: string[]) => void;
 
   // Computed values
   getContextPercentage: () => number;
@@ -296,6 +297,28 @@ export const useToolStore = create<ToolState>()(
       });
     },
 
+    restoreSessionUsage: (sessionId: string, usage: UsageData, processedIds?: string[]) => {
+      set((state) => {
+        // Pre-populate the session cache with usage from persisted data
+        // This is called when loading a conversation from disk
+        const existingCache = state.sessionCache[sessionId];
+        state.sessionCache[sessionId] = {
+          usage: { ...usage },
+          processedIds: processedIds ?? existingCache?.processedIds ?? [],
+          activeTools: existingCache?.activeTools ?? {},
+          completedTools: existingCache?.completedTools ?? [],
+        };
+
+        // If this is the current session, also update the active usage
+        if (state.currentSessionId === sessionId) {
+          state.sessionUsage = { ...usage };
+          if (processedIds) {
+            state.processedMessageIds = new Set(processedIds);
+          }
+        }
+      });
+    },
+
     getContextPercentage: () => {
       const state = get();
       const maxTokens = MODEL_CONTEXT_WINDOWS[state.model];
 
@@ -1316,22 +1316,62 @@ export const ConversationLoadingSchema = z
   })
   .strict();
 
+// Tool use schema for persisted messages (matches Rust ToolUseDto)
+const PersistedToolUseSchema = z
+  .object({
+    id: z.string(),
+    name: z.string(),
+    input: z.record(z.string(), z.unknown()),
+    output: z.string().optional(),
+    success: z.boolean().default(true),
+  })
+  .strict();
+
+// Token usage schema for persisted messages (matches Rust TokenUsageDto)
+const PersistedTokenUsageSchema = z
+  .object({
+    inputTokens: z.number(),
+    outputTokens: z.number(),
+    cacheReadInputTokens: z.number().optional(),
+    cacheCreationInputTokens: z.number().optional(),
+    totalCostUsd: z.number().optional(),
+  })
+  .strict();
+
+// Message schema for conversation:loaded (matches Rust MessageDto)
+// Uses transform to handle backwards compatibility with old data that may have
+// 'timestamp' instead of 'createdAt', or missing fields
+const PersistedMessageSchema = z
+  .object({
+    id: z.string(),
+    role: z.enum(['user', 'assistant', 'system']),
+    content: z.string(),
+    thinking: z.string().optional(),
+    // Support both old 'timestamp' and new 'createdAt' field names
+    createdAt: z.number().optional(),
+    timestamp: z.number().optional(),
+    toolUses: z.array(PersistedToolUseSchema).optional(),
+    usage: PersistedTokenUsageSchema.optional(),
+  })
+  .strip() // Remove extra fields from old data instead of rejecting
+  .transform((msg) => ({
+    id: msg.id,
+    role: msg.role,
+    content: msg.content,
+    thinking: msg.thinking,
+    // Prefer createdAt, fall back to timestamp, default to 0
+    createdAt: msg.createdAt ?? msg.timestamp ?? 0,
+    toolUses: msg.toolUses ?? [],
+    usage: msg.usage,
+  }));
+
 export const ConversationLoadedSchema = z
   .object({
     type: z.literal('conversation:loaded'),
     uuid: UUIDSchema,
     session_id: SessionIdSchema,
     title: z.string(),
-    messages: z.array(
-      z
-        .object({
-          id: z.string(),
-          role: z.enum(['user', 'assistant']),
-          content: z.string(),
-          timestamp: z.number(),
-        })
-        .strict()
-    ),
+    messages: z.array(PersistedMessageSchema),
   })
   .strict();
 
 
@@ -56,6 +56,9 @@ pub struct Message {
     /// Tool uses in this message
     #[serde(default, skip_serializing_if = "Vec::is_empty")]
     pub tool_uses: Vec<ToolUse>,
+    /// Token usage for this message (assistant messages only)
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub usage: Option<TokenUsage>,
 }
 
 /// A tool use within a message
@@ -76,6 +79,27 @@ pub struct ToolUse {
     pub success: bool,
 }
 
+/// Token usage statistics for a message
+#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct TokenUsage {
+    /// Input tokens consumed
+    #[serde(default)]
+    pub input_tokens: u32,
+    /// Output tokens generated
+    #[serde(default)]
+    pub output_tokens: u32,
+    /// Tokens read from cache
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub cache_read_input_tokens: Option<u32>,
+    /// Tokens used to create cache
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub cache_creation_input_tokens: Option<u32>,
+    /// Total cost in USD
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub total_cost_usd: Option<f64>,
+}
+
 // ============================================
 // Conversation Types
 // ============================================
@@ -641,6 +665,7 @@ mod tests {
             thinking: None,
             created_at: current_timestamp(),
             tool_uses: Vec::new(),
+            usage: None,
         }
     }