Fix chat message display issues with two-step processing

app.py

@@ -7,7 +7,7 @@ from datetime import datetime
 from pathlib import Path
 sys.path.append(str(Path(__file__).parent))
 
-# Import our new handler
+# Import modules
 from src.ui.chat_handler import chat_handler
 from utils.config import config
 from core.session import session_manager
@@ -22,7 +22,9 @@ import logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Initialize session state
+st.set_page_config(page_title="CosmicCat AI Assistant", page_icon="🐱", layout="wide")
+
+# Initialize session state properly
 if "messages" not in st.session_state:
     st.session_state.messages = []
 if "is_processing" not in st.session_state:
@@ -33,15 +35,16 @@ if "cosmic_mode" not in st.session_state:
     st.session_state.cosmic_mode = True
 if "show_welcome" not in st.session_state:
     st.session_state.show_welcome = True
-if "session_id" not in st.session_state:
-    st.session_state.session_id = f"sess_{int(time.time())}_{os.urandom(4).hex()}"
 if "last_processed_message" not in st.session_state:
     st.session_state.last_processed_message = ""
+if "session_id" not in st.session_state:
+    st.session_state.session_id = f"sess_{int(time.time())}_{abs(hash(str(time.time()))) % 10000}"
 
 # Start session tracking
-session_analytics.start_session_tracking("default_user", st.session_state.session_id)
-
-st.set_page_config(page_title="CosmicCat AI Assistant", page_icon="🐱", layout="wide")
+try:
+    session_analytics.start_session_tracking("default_user", st.session_state.session_id)
+except Exception as e:
+    logger.warning(f"Analytics session tracking failed: {e}")
 
 # Log page view
 session_analytics.track_interaction("default_user", st.session_state.session_id, "page_view", {
@@ -251,12 +254,12 @@ if st.session_state.show_welcome:
         st.markdown(greeting)
     st.session_state.show_welcome = False
 
-# Display conversation history
-for message in st.session_state.messages:
+# Display existing conversation history
+for message in st.session_state.get("messages", []):
     with st.chat_message(message["role"]):
         st.markdown(message["content"])
         if "timestamp" in message:
-            provider_info = f" (via {message.get('provider', 'unknown')})" if message["role"] == "assistant" else ""
+            provider_info = f" (via {message.get('provider', 'ollama')})" if message["role"] == "assistant" else ""
             st.caption(f"🕒 {message['timestamp']}{provider_info}")
 
 # Chat input with enhanced processing
@@ -264,12 +267,19 @@ user_input = st.chat_input("Type your message here...", key="chat_input")
 
 # Process message when received
 if user_input and user_input.strip():
-    # Prevent processing if already processing
-    if not getattr(chat_handler, 'is_processing', False):
+    # Handle user message display first
+    if not st.session_state.get('is_processing', False):
         chat_handler.process_user_message(user_input, selected_model_name)
     else:
         st.warning("Still processing your previous request...")
 
+# Handle AI response processing (triggered after user message display)
+if st.session_state.get('is_processing', False) and st.session_state.get('last_processed_message'):
+    chat_handler.process_ai_response(
+        st.session_state.last_processed_message, 
+        selected_model_name
+    )
+
 # About tab
 st.divider()
 tab1, = st.tabs(["ℹ️ About"])

src/ui/chat_handler.py

@@ -20,21 +20,22 @@ class ChatHandler:
             return
             
         # Prevent duplicate processing
-        if hasattr(st.session_state, 'last_processed_message'):
-            if st.session_state.last_processed_message == user_input:
-                return
-        
+        if st.session_state.get('last_processed_message') == user_input:
+            logger.info("Preventing duplicate message processing")
+            return
+            
+        # Set processing flag
+        st.session_state.is_processing = True
         st.session_state.last_processed_message = user_input
-        self.is_processing = True
         
         try:
-            # Show user message immediately
+            # Show user message immediately (this was missing!)
             timestamp = time.strftime("%H:%M:%S")
             with st.chat_message("user"):
                 st.markdown(user_input)
                 st.caption(f"🕒 {timestamp}")
             
-            # Add to session state history
+            # Add to session state history immediately
             if "messages" not in st.session_state:
                 st.session_state.messages = []
                 
@@ -44,6 +45,20 @@ class ChatHandler:
                 "timestamp": timestamp
             })
             
+            # Force immediate UI update
+            st.experimental_rerun()
+            
+        except Exception as e:
+            logger.error(f"Error in initial message display: {e}", exc_info=True)
+            st.session_state.is_processing = False
+            st.session_state.last_processed_message = ""
+    
+    def process_ai_response(self, user_input: str, selected_model: str):
+        """Process AI response after user message is displayed"""
+        if not user_input or not user_input.strip():
+            return
+            
+        try:
             # Show processing status
             with st.chat_message("assistant"):
                 status_placeholder = st.empty()
@@ -67,32 +82,54 @@ class ChatHandler:
                         response_placeholder.markdown(response)
                         
                         # Add to session history
+                        timestamp = time.strftime("%H:%M:%S")
                         st.session_state.messages.append({
                             "role": "assistant",
                             "content": response,
-                            "timestamp": time.strftime("%H:%M:%S"),
+                            "timestamp": timestamp,
                             "provider": provider_name
                         })
                     else:
                         status_placeholder.warning("⚠️ Empty response received")
                         response_placeholder.markdown("*No response generated. Please try again.*")
+                        timestamp = time.strftime("%H:%M:%S")
+                        st.session_state.messages.append({
+                            "role": "assistant",
+                            "content": "*No response generated. Please try again.*",
+                            "timestamp": timestamp,
+                            "provider": provider_name
+                        })
                         
                 except ProviderNotAvailableError as e:
                     status_placeholder.error("❌ No AI providers available")
                     response_placeholder.markdown("No AI providers are configured. Please check your settings.")
+                    timestamp = time.strftime("%H:%M:%S")
+                    st.session_state.messages.append({
+                        "role": "assistant",
+                        "content": "No AI providers are configured. Please check your settings.",
+                        "timestamp": timestamp
+                    })
                     logger.error(f"Provider not available: {e}")
                     
                 except Exception as e:
                     status_placeholder.error(f"❌ Error: {str(e)[:100]}...")
                     response_placeholder.markdown(f"Sorry, I encountered an error: {str(e)[:100]}...")
+                    timestamp = time.strftime("%H:%M:%S")
+                    st.session_state.messages.append({
+                        "role": "assistant",
+                        "content": f"Sorry, I encountered an error: {str(e)[:100]}...",
+                        "timestamp": timestamp
+                    })
                     logger.error(f"Chat processing error: {e}", exc_info=True)
                     
         except Exception as e:
-            logger.error(f"Unexpected error in process_user_message: {e}", exc_info=True)
+            logger.error(f"Unexpected error in process_ai_response: {e}", exc_info=True)
             st.error("An unexpected error occurred. Please try again.")
         finally:
-            self.is_processing = False
-            # Force UI refresh
+            # Clear processing flags
+            st.session_state.is_processing = False
+            st.session_state.last_processed_message = ""
+            time.sleep(0.1)
             st.experimental_rerun()
     
     def _get_provider_for_model(self, selected_model: str) -> str: