Improve the demo feature of the project

services/demo_session/app/services/cleanup_service.py

@@ -27,31 +27,55 @@ class DemoCleanupService:
     async def cleanup_expired_sessions(self) -> dict:
         """
         Find and cleanup all expired sessions
+        Also cleans up sessions stuck in PENDING for too long (>5 minutes)
 
         Returns:
             Cleanup statistics
         """
+        from datetime import timedelta
+
         logger.info("Starting demo session cleanup")
 
         now = datetime.now(timezone.utc)
+        stuck_threshold = now - timedelta(minutes=5)  # Sessions pending > 5 min are stuck
 
-        # Find expired sessions
+        # Find expired sessions (any status except EXPIRED and DESTROYED)
         result = await self.db.execute(
             select(DemoSession).where(
-                DemoSession.status == DemoSessionStatus.ACTIVE,
+                DemoSession.status.in_([
+                    DemoSessionStatus.PENDING,
+                    DemoSessionStatus.READY,
+                    DemoSessionStatus.PARTIAL,
+                    DemoSessionStatus.FAILED,
+                    DemoSessionStatus.ACTIVE  # Legacy status, kept for compatibility
+                ]),
                 DemoSession.expires_at < now
             )
         )
         expired_sessions = result.scalars().all()
 
+        # Also find sessions stuck in PENDING
+        stuck_result = await self.db.execute(
+            select(DemoSession).where(
+                DemoSession.status == DemoSessionStatus.PENDING,
+                DemoSession.created_at < stuck_threshold
+            )
+        )
+        stuck_sessions = stuck_result.scalars().all()
+
+        # Combine both lists
+        all_sessions_to_cleanup = list(expired_sessions) + list(stuck_sessions)
+
         stats = {
             "total_expired": len(expired_sessions),
+            "total_stuck": len(stuck_sessions),
+            "total_to_cleanup": len(all_sessions_to_cleanup),
             "cleaned_up": 0,
             "failed": 0,
             "errors": []
         }
 
-        for session in expired_sessions:
+        for session in all_sessions_to_cleanup:
             try:
                 # Mark as expired
                 session.status = DemoSessionStatus.EXPIRED
@@ -128,6 +152,11 @@ class DemoCleanupService:
 
         now = datetime.now(timezone.utc)
 
+        # Count by status
+        pending_count = len([s for s in all_sessions if s.status == DemoSessionStatus.PENDING])
+        ready_count = len([s for s in all_sessions if s.status == DemoSessionStatus.READY])
+        partial_count = len([s for s in all_sessions if s.status == DemoSessionStatus.PARTIAL])
+        failed_count = len([s for s in all_sessions if s.status == DemoSessionStatus.FAILED])
         active_count = len([s for s in all_sessions if s.status == DemoSessionStatus.ACTIVE])
         expired_count = len([s for s in all_sessions if s.status == DemoSessionStatus.EXPIRED])
         destroyed_count = len([s for s in all_sessions if s.status == DemoSessionStatus.DESTROYED])
@@ -135,13 +164,25 @@ class DemoCleanupService:
         # Find sessions that should be expired but aren't marked yet
         should_be_expired = len([
             s for s in all_sessions
-            if s.status == DemoSessionStatus.ACTIVE and s.expires_at < now
+            if s.status in [
+                DemoSessionStatus.PENDING,
+                DemoSessionStatus.READY,
+                DemoSessionStatus.PARTIAL,
+                DemoSessionStatus.FAILED,
+                DemoSessionStatus.ACTIVE
+            ] and s.expires_at < now
         ])
 
         return {
             "total_sessions": len(all_sessions),
-            "active_sessions": active_count,
-            "expired_sessions": expired_count,
-            "destroyed_sessions": destroyed_count,
+            "by_status": {
+                "pending": pending_count,
+                "ready": ready_count,
+                "partial": partial_count,
+                "failed": failed_count,
+                "active": active_count,  # Legacy
+                "expired": expired_count,
+                "destroyed": destroyed_count
+            },
             "pending_cleanup": should_be_expired
         }