Update app.py

app.py

@@ -212,7 +212,7 @@ def handle_chat(msg, chatbot, current_model, logged_in_username, user_chats_stat
         gr.Warning("⚠️ Load a model first!")
         return chatbot
 
-    # Cooldown check
+    # 3-second cooldown
     ok, new_time = enforce_cooldown(last_interaction_time.value)
     if not ok:
         return chatbot
@@ -224,24 +224,24 @@ def handle_chat(msg, chatbot, current_model, logged_in_username, user_chats_stat
     internal_history = chatbot[:] if chatbot else []
 
     try:
-        stream_gen = chat_stream(msg, internal_history, current_model, logged_in_username, user_chats_state, current_chat_title, mode_radio, auth_globals_state, file_upload, current_chat_context)
+        stream_gen = chat_stream(
+            msg, internal_history, current_model, logged_in_username, 
+            user_chats_state, current_chat_title, mode_radio, 
+            auth_globals_state, file_upload, current_chat_context
+        )
 
         for updated_history in stream_gen:
-            # Add approximate token count to last assistant message
-            if updated_history and updated_history[-1]["role"] == "assistant":
-                content = updated_history[-1]["content"]
-                tokens = len(content) // 4 + 50  # simple but good enough approximation
-                updated_history[-1]["tokens"] = tokens
-            yield updated_history
+            yield updated_history   # ← always yield to chatbot
 
     except Exception as e:
-        gr.Warning(f"⚠️ Error: {str(e)[:200]}")
+        gr.Warning(f"⚠️ Error: {str(e)[:250]}")
+        yield chatbot  # fallback
     finally:
         is_generating.value = False
         stop_btn.visible = False
 
-    # Update token display
-    return token_display.update(value=update_token_display(user_chats_state, current_chat_title))
+    # Update token count AFTER streaming finishes (using .then chain)
+    return chatbot  # important: return the chatbot itself
 
 
 # ================== GUI ==================