diff --git a/app/src/main/java/dev/kaizen/app/chat/ChatScreen.kt b/app/src/main/java/dev/kaizen/app/chat/ChatScreen.kt
index 1daca86..ad0f5ed 100644
--- a/app/src/main/java/dev/kaizen/app/chat/ChatScreen.kt
+++ b/app/src/main/java/dev/kaizen/app/chat/ChatScreen.kt
@@ -108,6 +108,7 @@ fun ChatScreen(
     var reasoningPreset by remember { mutableStateOf(ReasoningPreset.Standard) }
     var samplingPrefs by remember { mutableStateOf(SamplingPrefs()) }
     var usedTokens by remember { mutableStateOf(0) }
+    var chatModel by remember { mutableStateOf<String?>(null) }
 
     // Navigation State
     var currentScreen by remember { mutableStateOf(AppScreen.Chat) }
@@ -225,6 +226,7 @@ fun ChatScreen(
         conversationId = null
         viewingLockedChat = false
         usedTokens = 0
+        chatModel = null
     }
 
     fun unlockAndOpen(summary: ConversationSummary) {
@@ -314,6 +316,9 @@ fun ChatScreen(
             if (stored.isNotEmpty()) {
                 val entities = stored.mapIndexed { i, m -> m.toEntity(summary.id, i) }
                 chat.messageRepo.cacheMessages(summary.id, entities)
+                val lastAssistant = stored.lastOrNull { it.role == "assistant" }
+                lastAssistant?.model?.let { chatModel = it }
+                usedTokens = stored.sumOf { (it.inputTokens ?: 0) + (it.outputTokens ?: 0) }
                 if (stored.size != cachedList.size || stored.zip(cachedList).any { (s, c) -> s.id != c.id || s.content != c.content }) {
                     messages.clear()
                     stored.forEach { m ->
@@ -673,12 +678,13 @@ fun ChatScreen(
                                     }
                                 }
                                 Spacer(Modifier.width(8.dp))
+                                val displayModel = chatModel ?: session.config?.model ?: ""
                                 ModelPill(
-                                    label = modelDisplayName(session.config?.model ?: "", models),
+                                    label = modelDisplayName(displayModel, models),
                                     onClick = { haptics.tick(); showModelSheet = true },
                                 )
                                 Spacer(Modifier.weight(1f))
-                                val ctxLen = models.find { it.id == session.config?.model }?.context_length ?: 1_000_000
+                                val ctxLen = models.find { it.id == displayModel }?.context_length ?: 1_000_000
                                 TokenBadge(used = usedTokens, contextLength = ctxLen)
                             }
                         }
@@ -730,21 +736,19 @@ fun ChatScreen(
                                 )
                                 .then(if (isTablet) Modifier.widthIn(max = 720.dp) else Modifier)
                         ) {
-                            if (!imeVisible) {
-                                Spacer(Modifier.height(24.dp))
-                                ModePillsRow(
-                                    activeMode = activeMode,
-                                    onToggle = { mode ->
-                                        activeMode = if (activeMode == mode) null else mode
-                                        haptics.tick()
-                                    },
-                                    reasoningPreset = reasoningPreset,
-                                    onReasoningChange = { reasoningPreset = it; haptics.tick() },
-                                    samplingPrefs = samplingPrefs,
-                                    onSamplingChange = { samplingPrefs = it },
-                                )
-                                Spacer(Modifier.height(10.dp))
-                            }
+                            Spacer(Modifier.height(if (imeVisible) 6.dp else 24.dp))
+                            ModePillsRow(
+                                activeMode = activeMode,
+                                onToggle = { mode ->
+                                    activeMode = if (activeMode == mode) null else mode
+                                    haptics.tick()
+                                },
+                                reasoningPreset = reasoningPreset,
+                                onReasoningChange = { reasoningPreset = it; haptics.tick() },
+                                samplingPrefs = samplingPrefs,
+                                onSamplingChange = { samplingPrefs = it },
+                            )
+                            Spacer(Modifier.height(if (imeVisible) 6.dp else 10.dp))
                             ChatInput(
                                 value = input,
                                 onValueChange = { input = it },
@@ -791,7 +795,7 @@ fun ChatScreen(
                                 models = models,
                                 selectedModelId = session.config?.model ?: "",
                                 favorites = session.favorites,
-                                onSelectModel = { session.setModel(it); haptics.tick() },
+                                onSelectModel = { session.setModel(it); chatModel = it; haptics.tick() },
                                 onToggleFavorite = { session.toggleFavorite(it) },
                                 onDismiss = { showModelSheet = false },
                             )
diff --git a/app/src/main/java/dev/kaizen/app/net/KaizenApi.kt b/app/src/main/java/dev/kaizen/app/net/KaizenApi.kt
index 8a34ed5..4a2847d 100644
--- a/app/src/main/java/dev/kaizen/app/net/KaizenApi.kt
+++ b/app/src/main/java/dev/kaizen/app/net/KaizenApi.kt
@@ -87,6 +87,9 @@ data class StoredMessage(
     val attachments: List<Attachment> = emptyList(),
     val reasoning: String? = null,
     val sources: List<SearchSource>? = null,
+    val model: String? = null,
+    val inputTokens: Int? = null,
+    val outputTokens: Int? = null,
 )
 
 @Serializable private data class ConversationDetail(val messages: List<StoredMessage> = emptyList())