From 026e96b66c4c61e341b4bdc470ea09ea6765ceef Mon Sep 17 00:00:00 2001
From: Lukas <lukas.pupkalipinski@lpl-mind.de>
Date: Fri, 23 Jan 2026 15:18:12 +0100
Subject: [PATCH] Fix setup/loading flow and WebSocket connection

1. Setup redirect flow (setup -> loading -> login):
   - Add /loading to exempt paths
   - Redirect setup to login after completion
   - Redirect loading to login when initialization complete

2. Close pages after completion:
   - Block access to /setup after setup is done
   - Block access to /loading after initialization complete
   - Proper redirect handling prevents re-access

3. Fix WebSocket 403 error:
   - Change /ws/progress to /ws/connect (correct endpoint)
   - Add /ws/connect to exempt paths
   - Subscribe to 'system' room for progress updates
   - Fix message data handling format
---
 src/server/middleware/setup_redirect.py | 25 +++++++++++++++++++++++++
 src/server/web/templates/loading.html   | 20 ++++++++++++++------
 2 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/src/server/middleware/setup_redirect.py b/src/server/middleware/setup_redirect.py
index 5670396..d0ddb4b 100644
--- a/src/server/middleware/setup_redirect.py
+++ b/src/server/middleware/setup_redirect.py
@@ -32,9 +32,11 @@ class SetupRedirectMiddleware(BaseHTTPMiddleware):
     # Paths that should always be accessible, even without setup
     EXEMPT_PATHS = {
         "/setup",  # Setup page itself
+        "/loading",  # Loading page (initialization progress)
         "/login",  # Login page (needs to be accessible after setup)
         "/queue",  # Queue page (for initial load)
         "/api/auth/",  # All auth endpoints (setup, login, logout, register)
+        "/ws/connect",  # WebSocket connection (needed for loading page)
         "/api/queue/",  # Queue API endpoints
         "/api/downloads/",  # Download API endpoints
         "/api/config/",  # Config API (needed for setup and management)
@@ -115,6 +117,29 @@ class SetupRedirectMiddleware(BaseHTTPMiddleware):
         """
         path = request.url.path
         
+        # Check if trying to access setup or loading page after completion
+        if path in ("/setup", "/loading"):
+            if not self._needs_setup():
+                # Setup is complete, check loading status
+                if path == "/setup":
+                    # Redirect to loading if initialization is in progress
+                    # Otherwise redirect to login
+                    return RedirectResponse(url="/login", status_code=302)
+                elif path == "/loading":
+                    # Check if initialization is complete
+                    try:
+                        from src.server.database.connection import get_db_session
+                        from src.server.database.system_settings_service import SystemSettingsService
+                        
+                        async with get_db_session() as db:
+                            is_complete = await SystemSettingsService.is_initial_scan_completed(db)
+                            if is_complete:
+                                # Initialization complete, redirect to login
+                                return RedirectResponse(url="/login", status_code=302)
+                    except Exception:
+                        # If we can't check, allow access to loading page
+                        pass
+        
         # Skip setup check for exempt paths
         if self._is_path_exempt(path):
             return await call_next(request)
diff --git a/src/server/web/templates/loading.html b/src/server/web/templates/loading.html
index a08155b..e901978 100644
--- a/src/server/web/templates/loading.html
+++ b/src/server/web/templates/loading.html
@@ -294,13 +294,19 @@
 
         function connectWebSocket() {
             const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
-            const wsUrl = `${protocol}//${window.location.host}/ws/progress`;
+            const wsUrl = `${protocol}//${window.location.host}/ws/connect`;
 
             ws = new WebSocket(wsUrl);
 
             ws.onopen = () => {
                 console.log('WebSocket connected');
                 updateConnectionStatus(true);
+                
+                // Subscribe to system room for progress updates
+                ws.send(JSON.stringify({
+                    action: 'join',
+                    room: 'system'
+                }));
             };
 
             ws.onmessage = (event) => {
@@ -340,8 +346,10 @@
             }
         }
 
-        function handleProgressUpdate(data) {
-            const { type, status, title, message, percent, current, total, metadata } = data;
+        function handleProgressUpdate(message) {
+            // Handle WebSocket message format: { type: string, data: {...} }
+            const data = message.data || message;
+            const { type, status, title, message: msg, percent, current, total, metadata } = data;
 
             // Determine step ID based on type and metadata
             let stepId = metadata?.step_id || type;
@@ -351,7 +359,7 @@
                 createStep(stepId, title || stepTitles[stepId] || stepId);
             }
 
-            updateStep(stepId, status, message, percent, current, total);
+            updateStep(stepId, status, msg, percent, current, total);
 
             // Check for completion
             if (metadata?.initialization_complete) {
@@ -360,7 +368,7 @@
 
             // Handle errors
             if (status === 'failed') {
-                showError(message || 'An error occurred during initialization');
+                showError(msg || 'An error occurred during initialization');
             }
         }
 
@@ -474,7 +482,7 @@
         }
 
         function continueToApp() {
-            window.location.href = '/';
+            window.location.href = '/login';
         }
 
         // Start WebSocket connection when page loads