change app.py

app.py

@@ -23,7 +23,7 @@ SAMPLE_RATE = 16_000
 WARMUP_DURATION = 2.0  # seconds of silence for warmup
 WPM_WINDOW = 10  # seconds for running mean calculation
 CALIBRATION_PERIOD = 5  # seconds before showing WPM
-SESSION_TIMEOUT = int(os.environ.get("SESSION_TIMEOUT", "30"))      # Max 30s per session
+SESSION_TIMEOUT = int(os.environ.get("SESSION_TIMEOUT", "50"))      # Max 30s per session
 INACTIVITY_TIMEOUT = int(os.environ.get("INACTIVITY_TIMEOUT", "10"))  # Close after 10s silence
 MAX_CONCURRENT_SESSIONS = int(os.environ.get("MAX_SESSIONS", "50"))
 
@@ -44,7 +44,7 @@ _sessions_lock = threading.Lock()
 _session_registry = {}
 _registry_lock = threading.Lock()
 _last_cleanup = time.time()
-SESSION_REGISTRY_CLEANUP_INTERVAL = 60  # seconds
+SESSION_REGISTRY_CLEANUP_INTERVAL = 30  # seconds
 SESSION_MAX_AGE = 30  # 30 seconds - remove sessions older than this
 
 
@@ -80,20 +80,44 @@ def get_or_create_session(session_id: str = None) -> "UserSession":
 def _cleanup_stale_sessions():
     """Remove sessions that haven't been accessed recently."""
     now = time.time()
-    to_remove = []
+    to_remove_from_registry = []
+    to_remove_from_active = []
     
+    # Need both locks to safely check both dictionaries
     with _registry_lock:
-        for session_id, session in _session_registry.items():
-            last_accessed = getattr(session, '_last_accessed', 0)
-            # Remove if: not running AND (no activity for SESSION_MAX_AGE)
-            if not session.is_running and (now - last_accessed > SESSION_MAX_AGE):
-                to_remove.append(session_id)
-        
-        for session_id in to_remove:
-            _session_registry.pop(session_id, None)
+        with _sessions_lock:
+            # Find stale sessions in registry
+            for session_id, session in _session_registry.items():
+                # NEVER remove if still in active_sessions (websocket still running)
+                if session_id in _active_sessions:
+                    continue
+                    
+                last_accessed = getattr(session, '_last_accessed', 0)
+                # Remove if: not running AND not active AND old
+                if not session.is_running and (now - last_accessed > SESSION_MAX_AGE):
+                    to_remove_from_registry.append(session_id)
+            
+            # Find orphaned sessions in active_sessions (not in registry anymore)
+            for session_id, session in list(_active_sessions.items()):
+                if session_id not in _session_registry:
+                    # Orphaned - mark for removal
+                    if not session.is_running:
+                        to_remove_from_active.append(session_id)
+            
+            # Clean up registry
+            for session_id in to_remove_from_registry:
+                _session_registry.pop(session_id, None)
+            
+            # Clean up orphaned active sessions
+            for session_id in to_remove_from_active:
+                _active_sessions.pop(session_id, None)
+            
+            active_count = len(_active_sessions)
+            registry_count = len(_session_registry)
     
-    if to_remove:
-        print(f"Cleaned up {len(to_remove)} stale sessions. Active: {len(_session_registry)}")
+    total_cleaned = len(to_remove_from_registry) + len(to_remove_from_active)
+    if total_cleaned > 0:
+        print(f"Cleaned up {len(to_remove_from_registry)} stale + {len(to_remove_from_active)} orphaned sessions. Registry: {registry_count}, Active: {active_count}")
 
 
 def cleanup_session(session_id: str):