Bổ sung1 hour delay cho clear state và session clean do trigger on_chat_end

app.py

@@ -8,6 +8,8 @@ from typing import Dict, List, Any, Optional, Callable
 from dataclasses import dataclass, field
 import os
 import uuid
+import threading
+from datetime import datetime, timedelta
 
 
 API_BASE_URL = "https://sale-agent-m179.onrender.com"
@@ -23,6 +25,10 @@ class ConversationState:
     outputs: Optional[Dict[str, Any]] = None
     selected_model: str = "Gemini 2.0 Flash"
     product_model_search: bool = False
+    # New fields for delayed cleanup
+    pending_cleanup: bool = False
+    cleanup_timer: Optional[threading.Timer] = None
+    last_activity: datetime = field(default_factory=datetime.now)
     
     def reset(self):
         """Reset state to initial values"""
@@ -33,37 +39,118 @@ class ConversationState:
         self.outputs = None
         self.selected_model = "Gemini 2.0 Flash"
         self.product_model_search = False
+        # Reset cleanup fields but don't touch timers
+        self.pending_cleanup = False
+        self.last_activity = datetime.now()
+
+    def cancel_cleanup_timer(self):
+        """Cancel pending cleanup timer if exists"""
+        if self.cleanup_timer:
+            self.cleanup_timer.cancel()
+            self.cleanup_timer = None
+            print(f"🚫 Cancelled cleanup timer for session: {self.session_id}")
 
 
 class StateManager:
-    """Manages conversation state operations with per-session isolation"""
+    """Manages conversation state operations with per-session isolation and delayed cleanup"""
     
     # CLASS-LEVEL session storage for isolation between different browser sessions
     _session_states: Dict[str, ConversationState] = {}
+    _lock = threading.Lock()  # Thread safety for timer operations
     
     @staticmethod
     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}")
-        else:
-            print(f"🔄 Retrieved existing session state for: {session_id}")
-        
-        return StateManager._session_states[session_id]
+        with StateManager._lock:
+            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}")
+            else:
+                state = StateManager._session_states[session_id]
+                print(f"🔄 Retrieved existing session state for: {session_id}")
+                
+                # CRITICAL: If session was pending cleanup, cancel it because user is active again
+                if state.pending_cleanup:
+                    state.cancel_cleanup_timer()
+                    state.pending_cleanup = False
+                    print(f"♻️ User activity detected! Cancelled pending cleanup for: {session_id}")
+            
+            # Update activity timestamp
+            state.last_activity = datetime.now()
+            return state
+    
+    @staticmethod
+    def schedule_delayed_cleanup(session_id: str, delay_seconds: int = 3600):
+        """Schedule delayed cleanup for a session (default 1 hours for disconnect tolerance)"""
+        with StateManager._lock:
+            if session_id not in StateManager._session_states:
+                print(f"⚠️ Cannot schedule cleanup for non-existent session: {session_id}")
+                return
+                
+            state = StateManager._session_states[session_id]
+            
+            # Cancel existing timer if any
+            state.cancel_cleanup_timer()
+            
+            # Mark as pending cleanup
+            state.pending_cleanup = True
+            
+            # Schedule new cleanup
+            def delayed_cleanup():
+                print(f"⏰ Executing delayed cleanup for session: {session_id}")
+                StateManager._perform_actual_cleanup(session_id)
+            
+            state.cleanup_timer = threading.Timer(delay_seconds, delayed_cleanup)
+            state.cleanup_timer.start()
+            
+            print(f"⏱️ Scheduled cleanup in {delay_seconds}s for session: {session_id} (likely disconnect)")
     
     @staticmethod
-    def cleanup_session(session_id: str):
-        """Clean up session state when user disconnects"""
-        if session_id in StateManager._session_states:
+    def _perform_actual_cleanup(session_id: str):
+        """Perform the actual cleanup after delay"""
+        with StateManager._lock:
+            if session_id not in StateManager._session_states:
+                print(f"⚠️ Session already cleaned or doesn't exist: {session_id}")
+                return
+                
+            state = StateManager._session_states[session_id]
+            
+            # Double-check if session is still pending cleanup (user might have sent message)
+            if not state.pending_cleanup:
+                print(f"🚫 Cleanup cancelled - user activity detected for: {session_id}")
+                return
+            
+            # Perform API cleanup
+            try:
+                if API_BASE_URL:
+                    payload = {
+                        "reset_cache": True, 
+                        "reset_model": False, 
+                        "session_id": session_id
+                    }
+                    response = requests.post(f"{API_BASE_URL}/clear_memory", json=payload, timeout=30)
+                    print(f"Clear memory response for {session_id}: {response.status_code}")
+            except Exception as e:
+                print(f"Warning: clear_memory failed for {session_id}: {e}")
+            
+            # Remove from memory
             del StateManager._session_states[session_id]
-            print(f"🗑️ Cleaned up session state for: {session_id}")
+            print(f"🗑️ Successfully cleaned up session: {session_id}")
+    
+    @staticmethod
+    def cleanup_session_immediate(session_id: str):
+        """Immediate cleanup (for testing or forced cleanup)"""
+        with StateManager._lock:
+            if session_id in StateManager._session_states:
+                state = StateManager._session_states[session_id]
+                state.cancel_cleanup_timer()
+                StateManager._perform_actual_cleanup(session_id)
     
     @staticmethod
     async def clear_chat_state(state: ConversationState):
-        """Clear all conversation history and reset state via API"""
+        """Clear all conversation history and reset state via API (but keep session alive)"""
         if state.session_id is not None and API_BASE_URL:
             try:
                 payload = {
@@ -71,12 +158,12 @@ class StateManager:
                     "reset_model": False, 
                     "session_id": state.session_id
                 }
-                response = requests.post(f"{API_BASE_URL}/clear_memory", json=payload)
+                response = requests.post(f"{API_BASE_URL}/clear_memory", json=payload, timeout=30)
                 print(f"Clear memory response: {response.status_code}")
             except Exception as e:
                 print(f"Warning: clear_memory failed: {e}")
         
-        # Reset state but keep session_id
+        # Reset state but keep session_id and don't trigger cleanup
         session_id = state.session_id
         state.reset()
         state.session_id = session_id
@@ -85,11 +172,28 @@ class StateManager:
     def change_model(state: ConversationState, model_name: str):
         """Change the selected model"""
         state.selected_model = model_name
+        state.last_activity = datetime.now()
     
     @staticmethod
     def toggle_product_model_search(state: ConversationState):
         """Toggle product model search mode"""
         state.product_model_search = not state.product_model_search
+        state.last_activity = datetime.now()
+    
+    @staticmethod
+    def get_session_status() -> Dict[str, Dict[str, Any]]:
+        """Get status of all sessions (for debugging)"""
+        with StateManager._lock:
+            status = {}
+            for session_id, state in StateManager._session_states.items():
+                status[session_id] = {
+                    "pending_cleanup": state.pending_cleanup,
+                    "has_timer": state.cleanup_timer is not None,
+                    "last_activity": state.last_activity.isoformat(),
+                    "selected_model": state.selected_model,
+                    "product_model_search": state.product_model_search
+                }
+            return status
 
 
 class ChatService:
@@ -103,6 +207,10 @@ class ChatService:
     ) -> 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}")
+        
+        # Update activity timestamp - this is KEY to prevent cleanup during active use
+        state.last_activity = datetime.now()
+        
         start = time.perf_counter()
         
         if not API_BASE_URL:
@@ -441,10 +549,8 @@ def ensure_session_state() -> Optional[ConversationState]:
         session_id = cl.user_session.get("session_id")
         
         if not 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"Lỗi: Không lấy được session id ở ensure_session_state")
+            return None
         
         return StateManager.get_or_create_session_state(session_id)
         
@@ -457,7 +563,7 @@ def get_current_session_state() -> Optional[ConversationState]:
     """Get current session state using Chainlit's session system"""
     try:
         # Use Chainlit's user session to get unique session ID
-        chainlit_session_id = cl.user_session.get("session_id")  # Fixed: added key parameter
+        chainlit_session_id = cl.user_session.get("session_id")
         
         if chainlit_session_id:
             return StateManager.get_or_create_session_state(chainlit_session_id)
@@ -476,6 +582,9 @@ async def on_chat_start():
     if not session_id:
         session_id = str(uuid.uuid4())
         cl.user_session.set("session_id", session_id)
+        print(f"🆕 Generated new session_id: {session_id}")
+    else:
+        print(f"🔄 Reusing existing session_id: {session_id}")
         
     app_state = StateManager.get_or_create_session_state(session_id)
 
@@ -498,21 +607,21 @@ async def on_chat_start():
 
 @cl.on_chat_end
 async def on_chat_end():
-    """Clean up when chat session ends (user closes tab, new chat, exit UI)"""
+    """Handle chat session end with delayed cleanup mechanism"""
     try:
         session_id = cl.user_session.get("session_id")
-        print(f"on_chat_end triggered for session {session_id}")
+        print(f"📤 on_chat_end triggered for session {session_id}")
         
         if session_id:
-            app_state = StateManager.get_or_create_session_state(session_id)
-            await StateManager.clear_chat_state(app_state)
-            StateManager.cleanup_session(session_id)
-            print(f"✅ Properly cleaned session {session_id}")
+            # Schedule delayed cleanup instead of immediate cleanup
+            # Use shorter delay (30s) since this is likely just a temporary disconnect
+            StateManager.schedule_delayed_cleanup(session_id, delay_seconds=3600)
+            print(f"⏳ Scheduled delayed cleanup for session {session_id} (1h delay for disconnect tolerance)")
         else:
-            print("⚠️ No session_id found in on_chat_end (maybe reconnect case)")
+            print("⚠️ No session_id found in on_chat_end")
     except Exception as e:
-        print(f"⚠️ Error during cleanup: {e}")
-    
+        print(f"⚠️ Error during on_chat_end: {e}")
+
 
 # ACTION CALLBACKS - All use ensure_session_state() for better reliability
 @cl.action_callback("show_specs")
@@ -657,6 +766,30 @@ async def on_select_model_2(action):
     await UIService.send_message_with_buttons("✅ Đã chuyển sang **Gemini 2.0 Flash Lite**", app_state, author="assistant")
 
 
+# DEBUG ENDPOINTS (optional - for monitoring session status)
+@cl.action_callback("debug_sessions")
+async def on_debug_sessions(action):
+    """Debug action to show session status (can be added to debug builds)"""
+    try:
+        status = StateManager.get_session_status()
+        debug_content = "🔍 **Debug: Session Status**\n\n"
+        
+        if not status:
+            debug_content += "No active sessions."
+        else:
+            for session_id, info in status.items():
+                debug_content += f"**Session: {session_id[:8]}...**\n"
+                debug_content += f"- Pending cleanup: {info['pending_cleanup']}\n"
+                debug_content += f"- Has timer: {info['has_timer']}\n"
+                debug_content += f"- Last activity: {info['last_activity']}\n"
+                debug_content += f"- Model: {info['selected_model']}\n"
+                debug_content += f"- Product search: {info['product_model_search']}\n\n"
+        
+        await cl.Message(content=debug_content, author="assistant").send()
+    except Exception as e:
+        await cl.Message(content=f"Debug error: {e}", author="assistant").send()
+
+
 @cl.on_message
 async def main(message: cl.Message):
     """Main message handler"""
@@ -683,4 +816,7 @@ async def main(message: cl.Message):
     typing_msg.content = response
     typing_msg.actions = UIService.create_action_buttons(app_state)
     typing_msg.author = "assistant"
-    await typing_msg.update()
+    await typing_msg.update()
+
+
+