Thử nghiệm chống trigger on_chat_end khi websocket disconnect

app.py

@@ -4,14 +4,13 @@ import chainlit as cl
 import pandas as pd
 import requests
 import asyncio
-import threading
 from typing import Dict, List, Any, Optional, Callable
 from dataclasses import dataclass, field
 import os
 import uuid
 
 
-API_BASE_URL = "https://sale-agent-m179.onrender.com"
+API_BASE_URL = os.getenv("API_BASE_URL")
 
 
 @dataclass
@@ -37,97 +36,34 @@ class ConversationState:
 
 
 class StateManager:
-    """Thread-safe StateManager with duplicate prevention"""
+    """Manages conversation state operations with per-session isolation"""
     
+    # CLASS-LEVEL session storage for isolation between different browser sessions
     _session_states: Dict[str, ConversationState] = {}
-    _cleanup_flags: Dict[str, bool] = {}
-    _creation_locks: Dict[str, threading.Lock] = {}
-    _cleanup_locks: Dict[str, threading.Lock] = {}
-    _global_lock = threading.Lock()
     
     @staticmethod
-    def get_session_state(session_id: str) -> Optional[ConversationState]:
-        """Thread-safe get session state"""
-        if session_id in StateManager._cleanup_flags:
-            print(f"⚠️ Session {session_id[:8]}... is being cleaned up, returning None")
-            return None
-            
-        return StateManager._session_states.get(session_id)
-    
-    @staticmethod
-    def create_session_state(session_id: str) -> ConversationState:
-        """Thread-safe create session state with duplicate prevention"""
-        # Get or create lock for this session
-        with StateManager._global_lock:
-            if session_id not in StateManager._creation_locks:
-                StateManager._creation_locks[session_id] = threading.Lock()
-            lock = StateManager._creation_locks[session_id]
-        
-        # Use session-specific lock to prevent duplicate creation
-        with lock:
-            # Double-check if session already exists
-            if session_id in StateManager._session_states:
-                print(f"🔄 Session {session_id[:8]}... already exists, returning existing")
-                return StateManager._session_states[session_id]
-            
-            # Check if being cleaned up
-            if session_id in StateManager._cleanup_flags:
-                print(f"⚠️ Session {session_id[:8]}... is being cleaned up, waiting...")
-                # Wait a bit and retry
-                time.sleep(0.1)
-                return StateManager.create_session_state(session_id)
-            
-            # Create new session
+    def get_or_create_session_state(session_id: str) -> ConversationState:
+        """Get existing session state or create new one"""
+        if session_id not in StateManager._session_states:
             state = ConversationState()
             state.session_id = session_id
             StateManager._session_states[session_id] = state
-            print(f"🆕 Created new session state for: {session_id[:8]}...")
-            return state
+            print(f"🆕 Created new session state for: {session_id}")
+        else:
+            print(f"🔄 Retrieved existing session state for: {session_id}")
+        
+        return StateManager._session_states[session_id]
     
     @staticmethod
     def cleanup_session(session_id: str):
-        """Thread-safe cleanup with duplicate prevention"""
-        # Get or create lock for this session
-        with StateManager._global_lock:
-            if session_id not in StateManager._cleanup_locks:
-                StateManager._cleanup_locks[session_id] = threading.Lock()
-            lock = StateManager._cleanup_locks[session_id]
-        
-        # Use session-specific lock to prevent duplicate cleanup
-        with lock:
-            # Check if already being cleaned up
-            if session_id in StateManager._cleanup_flags:
-                print(f"⚠️ Session {session_id[:8]}... already being cleaned up, skipping")
-                return
-            
-            # Mark as being cleaned up
-            StateManager._cleanup_flags[session_id] = True
-            
-            try:
-                # Cleanup session data
-                if session_id in StateManager._session_states:
-                    del StateManager._session_states[session_id]
-                    print(f"🗑️ Cleaned up session state for: {session_id[:8]}...")
-                else:
-                    print(f"⚠️ No session state found for cleanup: {session_id[:8]}...")
-                
-                # Cleanup locks
-                if session_id in StateManager._creation_locks:
-                    del StateManager._creation_locks[session_id]
-                
-            finally:
-                # Always remove cleanup flag
-                if session_id in StateManager._cleanup_flags:
-                    del StateManager._cleanup_flags[session_id]
-                    
-                # Remove cleanup lock
-                with StateManager._global_lock:
-                    if session_id in StateManager._cleanup_locks:
-                        del StateManager._cleanup_locks[session_id]
+        """Clean up session state when user disconnects"""
+        if session_id in StateManager._session_states:
+            del StateManager._session_states[session_id]
+            print(f"🗑️ Cleaned up session state for: {session_id}")
     
     @staticmethod
     async def clear_chat_state(state: ConversationState):
-        """Clear chat state with better error handling"""
+        """Clear all conversation history and reset state via API"""
         if state.session_id is not None and API_BASE_URL:
             try:
                 payload = {
@@ -135,10 +71,10 @@ class StateManager:
                     "reset_model": False, 
                     "session_id": state.session_id
                 }
-                response = requests.post(f"{API_BASE_URL}/clear_memory", json=payload, timeout=30)
-                print(f"Clear memory response: {response.status_code} for session: {state.session_id[:8]}...")
+                response = requests.post(f"{API_BASE_URL}/clear_memory", json=payload)
+                print(f"Clear memory response: {response.status_code}")
             except Exception as e:
-                print(f"Warning: clear_memory failed for session {state.session_id[:8]}...: {e}")
+                print(f"Warning: clear_memory failed: {e}")
         
         # Reset state but keep session_id
         session_id = state.session_id
@@ -166,7 +102,7 @@ class ChatService:
         image_path: Optional[str] = None
     ) -> str:
         """Handle chat responses with image support"""
-        print(f"🔄 === DEBUG STATE ===\n Chat request with model: {state.selected_model}, Product Model Search: {state.product_model_search}, Session ID: {state.session_id[:8] if state.session_id else 'None'}...")
+        print(f"🔄 === DEBUG STATE ===\n Chat request with model: {state.selected_model}, Product Model Search: {state.product_model_search}, Session ID: {state.session_id}")
         start = time.perf_counter()
         
         if not API_BASE_URL:
@@ -498,134 +434,92 @@ class UIService:
         return msg
 
 
-# HELPER FUNCTIONS - Improved session management with duplicate prevention
-_message_processing = set()  # Track processing messages to prevent duplicates
-
-def get_current_session_state() -> Optional[ConversationState]:
-    """Get current session state with better error handling"""
+# HELPER FUNCTIONS: Session management with proper error handling
+def ensure_session_state() -> Optional[ConversationState]:
+    """Ensure session state exists, create if not"""
     try:
         session_id = cl.user_session.get("session_id")
+        
         if not session_id:
-            print("⚠️ No session ID found in Chainlit session")
-            return None
-            
-        return StateManager.get_session_state(session_id)
+            # Chỉ tạo mới khi chưa có (tức là lần đầu mở tab)
+            session_id = str(uuid.uuid4())
+            cl.user_session.set("session_id", session_id)
+            print(f"🆕 Created new session_id for new chat: {session_id}")
+        
+        return StateManager.get_or_create_session_state(session_id)
+        
     except Exception as e:
-        print(f"⚠️ Error getting session state: {e}")
+        print(f"⚠️ Error ensuring session state: {e}")
         return None
 
 
-def ensure_session_state() -> Optional[ConversationState]:
-    """Get existing session state or create new one if needed"""
+def get_current_session_state() -> Optional[ConversationState]:
+    """Get current session state using Chainlit's session system"""
     try:
-        session_id = cl.user_session.get("session_id")
-        if not session_id:
-            # Create new session ID
-            session_id = str(uuid.uuid4())
-            cl.user_session.set("session_id", session_id)
-            print(f"🔄 Created new session ID: {session_id[:8]}...")
+        # Use Chainlit's user session to get unique session ID
+        chainlit_session_id = cl.user_session.get("session_id")  # Fixed: added key parameter
         
-        # Try to get existing session first
-        state = StateManager.get_session_state(session_id)
-        if state is None:
-            # Create new session if not exists
-            state = StateManager.create_session_state(session_id)
-            
-        return state
+        if chainlit_session_id:
+            return StateManager.get_or_create_session_state(chainlit_session_id)
+        else:
+            print("⚠️ No Chainlit session ID found")
+            return None
     except Exception as e:
-        print(f"⚠️ Error ensuring session state: {e}")
+        print(f"⚠️ Error getting session state: {e}")
         return None
 
 
 @cl.on_chat_start
 async def on_chat_start():
-    """Initialize the chat session with duplicate prevention"""
-    try:
-        # Check if session already exists
-        existing_session_id = cl.user_session.get("session_id")
-        if existing_session_id:
-            existing_state = StateManager.get_session_state(existing_session_id)
-            if existing_state:
-                print(f"🔄 Reusing existing session: {existing_session_id[:8]}...")
-                session_id = existing_session_id
-                app_state = existing_state
-            else:
-                # Session ID exists but state doesn't, create new state
-                print(f"🔄 Recreating state for existing session: {existing_session_id[:8]}...")
-                app_state = StateManager.create_session_state(existing_session_id)
-                session_id = existing_session_id
-        else:
-            # Generate completely new session
-            session_id = str(uuid.uuid4())
-            cl.user_session.set("session_id", session_id)
-            app_state = StateManager.create_session_state(session_id)
-        
-        print(f"📋 Session initialized: {session_id[:8]}...")
-        
-        await cl.Message(
-            content=f"🛍️ **RangDong Sales Agent** (Session: {session_id[:8]}...)\n\n"
+    """Initialize the chat session"""
+    session_id = cl.user_session.get("session_id")
+    if not session_id:
+        session_id = str(uuid.uuid4())
+        cl.user_session.set("session_id", session_id)
+
+    app_state = StateManager.get_or_create_session_state(session_id)
+
+    await cl.Message(
+        content=f"🛍️ **RangDong Sales Agent** (Session: {session_id[:8]}...)\n\n"
                    f"Xin chào! Tôi có thể giúp bạn tìm kiếm và tư vấn sản phẩm RangDong. Hãy thử các câu hỏi mẫu:\n\n"
                    f"- Tìm sản phẩm bình giữ nhiệt dung tích dưới 2 lít\n"
                    f"- Tìm sản phẩm ổ cắm thông minh\n"
                    f"- Tư vấn cho tôi đèn học chống cận cho con gái của tôi học lớp 6",
-            author="assistant"
-        ).send()
-        
-        # Create initial action buttons
-        actions = UIService.create_start_buttons(app_state)
-        await cl.Message(
-            content="Sử dụng nút bên dưới để cấu hình:", 
-            actions=actions,
-            author="assistant"
-        ).send()
-        
-    except Exception as e:
-        print(f"⚠️ Error in on_chat_start: {e}")
-        await cl.Message(
-            content="❌ Lỗi khởi tạo session. Vui lòng refresh trang.",
-            author="assistant"
-        ).send()
+        author="assistant"
+    ).send()
+
+    actions = UIService.create_start_buttons(app_state)
+    await cl.Message(
+        content="Sử dụng nút bên dưới để cấu hình:", 
+        actions=actions,
+        author="assistant"
+    ).send()
 
 
 @cl.on_chat_end
 async def on_chat_end():
-    """Clean up when chat session ends with duplicate prevention"""
+    """Clean up when chat session ends (user closes tab, new chat, exit UI)"""
     try:
         session_id = cl.user_session.get("session_id")
-        if not session_id:
-            print("⚠️ No session ID found during cleanup")
-            return
-        
-        print(f"🔚 Starting cleanup for session: {session_id[:8]}...")
         
-        # Get existing session without creating new one
-        app_state = StateManager.get_session_state(session_id)
-        
-        if app_state:
-            # Clear chat state via API
+        if session_id:
+            app_state = StateManager.get_or_create_session_state(session_id)
             await StateManager.clear_chat_state(app_state)
-            print(f"✅ Chat state cleared for session: {session_id[:8]}...")
+            StateManager.cleanup_session(session_id)
+            print(f"✅ Properly cleaned session {session_id}")
         else:
-            print(f"⚠️ No app state found for session: {session_id[:8]}...")
-        
-        # Always cleanup session from memory
-        StateManager.cleanup_session(session_id)
-        print(f"✅ Session cleanup completed for: {session_id[:8]}...")
-        
+            print("⚠️ No session_id found in on_chat_end (maybe reconnect case)")
     except Exception as e:
         print(f"⚠️ Error during cleanup: {e}")
+    
 
-
-# ACTION CALLBACKS - Improved error handling
+# ACTION CALLBACKS - All use ensure_session_state() for better reliability
 @cl.action_callback("show_specs")
 async def on_show_specs(action):
     """Handle show specifications action"""
-    app_state = get_current_session_state()
+    app_state = ensure_session_state()
     if app_state is None:
-        await cl.Message(
-            content="❌ Session không tồn tại hoặc đã bị đóng. Vui lòng refresh trang để tạo session mới.",
-            author="assistant"
-        ).send()
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
         return
         
     specs_content = DisplayService.show_specs(app_state)
@@ -635,12 +529,9 @@ async def on_show_specs(action):
 @cl.action_callback("show_advantages")
 async def on_show_advantages(action):
     """Handle show advantages action"""
-    app_state = get_current_session_state()
+    app_state = ensure_session_state()
     if app_state is None:
-        await cl.Message(
-            content="❌ Session không tồn tại hoặc đã bị đóng. Vui lòng refresh trang để tạo session mới.",
-            author="assistant"
-        ).send()
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
         return
         
     adv_content = DisplayService.show_advantages(app_state)
@@ -650,42 +541,31 @@ async def on_show_advantages(action):
 @cl.action_callback("show_packages")
 async def on_show_packages(action):
     """Handle show packages action"""
-    app_state = get_current_session_state()
+    app_state = ensure_session_state()
     if app_state is None:
-        await cl.Message(
-            content="❌ Session không tồn tại hoặc đã bị đóng. Vui lòng refresh trang để tạo session mới.",
-            author="assistant"
-        ).send()
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
         return
         
     pkg_content = DisplayService.show_solution_packages(app_state)
     await UIService.send_message_with_buttons(pkg_content, app_state, author="assistant")
 
-
 @cl.action_callback("show_all_products")
 async def on_show_all_products(action):
     """Handle show all products action"""
-    app_state = get_current_session_state()
+    app_state = ensure_session_state()
     if app_state is None:
-        await cl.Message(
-            content="❌ Session không tồn tại hoặc đã bị đóng. Vui lòng refresh trang để tạo session mới.",
-            author="assistant"
-        ).send()
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
         return
         
     all_products_content = DisplayService.show_all_products_table(app_state)
     await UIService.send_message_with_buttons(all_products_content, app_state, author="assistant")
 
-
 @cl.action_callback("toggle_product_search")
 async def on_toggle_product_search(action):
     """Handle toggle product model search action"""
-    app_state = get_current_session_state()
+    app_state = ensure_session_state()
     if app_state is None:
-        await cl.Message(
-            content="❌ Session không tồn tại hoặc đã bị đóng. Vui lòng refresh trang để tạo session mới.",
-            author="assistant"
-        ).send()
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
         return
         
     StateManager.toggle_product_model_search(app_state)
@@ -704,12 +584,9 @@ async def on_toggle_product_search(action):
 @cl.action_callback("change_model")
 async def on_change_model(action):
     """Handle model change action"""
-    app_state = get_current_session_state()
+    app_state = ensure_session_state()
     if app_state is None:
-        await cl.Message(
-            content="❌ Session không tồn tại hoặc đã bị đóng. Vui lòng refresh trang để tạo session mới.",
-            author="assistant"
-        ).send()
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
         return
         
     models = ["Gemini 2.0 Flash", "Gemini 2.5 Flash Lite", "Gemini 2.0 Flash Lite"]
@@ -733,12 +610,9 @@ async def on_change_model(action):
 @cl.action_callback("back_to_main")
 async def on_back_to_main(action):
     """Handle back to main menu action"""
-    app_state = get_current_session_state()
+    app_state = ensure_session_state()
     if app_state is None:
-        await cl.Message(
-            content="❌ Session không tồn tại hoặc đã bị đóng. Vui lòng refresh trang để tạo session mới.",
-            author="assistant"
-        ).send()
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
         return
         
     actions = UIService.create_action_buttons(app_state)
@@ -751,12 +625,9 @@ async def on_back_to_main(action):
 
 @cl.action_callback("select_model_0")
 async def on_select_model_0(action):
-    app_state = get_current_session_state()
+    app_state = ensure_session_state()
     if app_state is None:
-        await cl.Message(
-            content="❌ Session không tồn tại hoặc đã bị đóng. Vui lòng refresh trang để tạo session mới.",
-            author="assistant"
-        ).send()
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
         return
         
     StateManager.change_model(app_state, "Gemini 2.0 Flash")
@@ -765,12 +636,9 @@ async def on_select_model_0(action):
 
 @cl.action_callback("select_model_1")
 async def on_select_model_1(action):
-    app_state = get_current_session_state()
+    app_state = ensure_session_state()
     if app_state is None:
-        await cl.Message(
-            content="❌ Session không tồn tại hoặc đã bị đóng. Vui lòng refresh trang để tạo session mới.",
-            author="assistant"
-        ).send()
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
         return
         
     StateManager.change_model(app_state, "Gemini 2.5 Flash Lite")
@@ -779,12 +647,9 @@ async def on_select_model_1(action):
 
 @cl.action_callback("select_model_2")
 async def on_select_model_2(action):
-    app_state = get_current_session_state()
+    app_state = ensure_session_state()
     if app_state is None:
-        await cl.Message(
-            content="❌ Session không tồn tại hoặc đã bị đóng. Vui lòng refresh trang để tạo session mới.",
-            author="assistant"
-        ).send()
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
         return
         
     StateManager.change_model(app_state, "Gemini 2.0 Flash Lite")
@@ -793,55 +658,28 @@ async def on_select_model_2(action):
 
 @cl.on_message
 async def main(message: cl.Message):
-    """Main message handler with duplicate prevention"""
-    try:
-        # Create unique message ID to prevent duplicate processing
-        message_id = f"{message.content[:50]}_{hash(message.content)}_{time.time()}"
-        
-        if message_id in _message_processing:
-            print(f"⚠️ Message already being processed: {message_id}")
-            return
-        
-        _message_processing.add(message_id)
-        
-        try:
-            app_state = ensure_session_state()
-            if app_state is None:
-                await cl.Message(
-                    content="❌ Session không tồn tại hoặc đã bị đóng. Vui lòng refresh trang để tạo session mới.",
-                    author="assistant"
-                ).send()
-                return
-            
-            print(f"📝 Processing message for session: {app_state.session_id[:8]}...")
-            
-            # Handle images if present
-            image_path = None
-            if message.elements:
-                for element in message.elements:
-                    if isinstance(element, cl.Image):
-                        image_path = element.path
-                        break
-            
-            # Show typing animation
-            typing_msg = await UIService.create_typing_animation()
-            
-            # Get response from API
-            response = await ChatService.respond_to_chat(app_state, message.content, image_path)
-            
-            # Update the typing message with final response and buttons
-            typing_msg.content = response
-            typing_msg.actions = UIService.create_action_buttons(app_state)
-            typing_msg.author = "assistant"
-            await typing_msg.update()
-            
-        finally:
-            # Always remove from processing set
-            _message_processing.discard(message_id)
-            
-    except Exception as e:
-        print(f"⚠️ Error in message handler: {e}")
-        await cl.Message(
-            content="❌ Lỗi xử lý tin nhắn. Vui lòng thử lại.",
-            author="assistant"
-        ).send()
+    """Main message handler"""
+    app_state = ensure_session_state()
+    if app_state is None:
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
+        return
+    
+    # Handle images if present
+    image_path = None
+    if message.elements:
+        for element in message.elements:
+            if isinstance(element, cl.Image):
+                image_path = element.path
+                break
+    
+    # Show typing animation
+    typing_msg = await UIService.create_typing_animation()
+    
+    # Get response from API
+    response = await ChatService.respond_to_chat(app_state, message.content, image_path)
+    
+    # Update the typing message with final response and buttons
+    typing_msg.content = response
+    typing_msg.actions = UIService.create_action_buttons(app_state)
+    typing_msg.author = "assistant"
+    await typing_msg.update()