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 4306b07..886badc 100644
--- a/app/src/main/java/dev/kaizen/app/chat/ChatScreen.kt
+++ b/app/src/main/java/dev/kaizen/app/chat/ChatScreen.kt
@@ -639,8 +639,9 @@ fun ChatScreen(
                                     modifier = Modifier.weight(1f, fill = false),
                                 )
                                 if (usedTokens > 0 || messages.isNotEmpty()) {
+                                    val ctxLen = models.find { it.id == session.config?.model }?.context_length ?: 1_000_000
                                     Spacer(Modifier.width(8.dp))
-                                    TokenBadge(used = usedTokens, contextLength = 1_000_000)
+                                    TokenBadge(used = usedTokens, contextLength = ctxLen)
                                 }
                             }
                         }
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 71739ab..224ceb2 100644
--- a/app/src/main/java/dev/kaizen/app/net/KaizenApi.kt
+++ b/app/src/main/java/dev/kaizen/app/net/KaizenApi.kt
@@ -47,6 +47,7 @@ data class KaizenModel(
     val id: String,
     val name: String? = null,
     val provider: String? = null,
+    val context_length: Int? = null,
 )
 
 @Serializable private data class ModelsResponse(val models: List<KaizenModel> = emptyList())