fix(indexing): WebSocket race condition, NaN stats, file streaming (#154)

backend/routes/ws_playground.py

@@ -152,6 +152,9 @@ async def websocket_playground_index(websocket: WebSocket, job_id: str):
                         job_id=job_id[:12],
                         event_type=event_type
                     )
+                    # Small delay to ensure client processes message before close
+                    # This prevents race condition where onclose fires before onmessage
+                    await asyncio.sleep(0.2)
                     break
 
             except json.JSONDecodeError:

backend/services/anonymous_indexer.py

@@ -43,9 +43,9 @@ def to_dict(self) -> dict:
 @dataclass
 class JobStats:
     """Final stats for completed job."""
-    files_indexed: int = 0
-    functions_found: int = 0
-    time_taken_seconds: float = 0
+    files_processed: int = 0
+    functions_indexed: int = 0
+    indexing_time_seconds: float = 0
 
     def to_dict(self) -> dict:
         return asdict(self)
@@ -345,12 +345,18 @@ async def run_indexing_job(
         job_manager.update_status(job_id, JobStatus.PROCESSING)
 
         # Progress callback for real-time updates
-        async def progress_callback(files_processed: int, functions_found: int, total: int):
+        async def progress_callback(
+            files_processed: int,
+            functions_found: int,
+            total: int,
+            current_file: Optional[str] = None
+        ):
             job_manager.update_progress(
                 job_id,
                 files_processed=files_processed,
                 functions_found=functions_found,
-                files_total=total
+                files_total=total,
+                current_file=current_file
             )
 
         # Run indexing with timeout
@@ -370,9 +376,9 @@ async def progress_callback(files_processed: int, functions_found: int, total: i
         # --- Step 3: Mark complete ---
         elapsed = time.time() - start_time
         stats = JobStats(
-            files_indexed=file_count,
-            functions_found=total_functions,
-            time_taken_seconds=round(elapsed, 2)
+            files_processed=file_count,
+            functions_indexed=total_functions,
+            indexing_time_seconds=round(elapsed, 2)
         )
 
         job_manager.update_status(

backend/services/indexer_optimized.py

@@ -672,8 +672,11 @@ async def index_repository_with_progress(
 
             files_processed = min(i + self.FILE_BATCH_SIZE, total_files)
 
-            # Send progress update
-            await progress_callback(files_processed, len(all_functions_data), total_files)
+            # Get the last file in this batch for display
+            current_file = batch[-1].name if batch else None
+
+            # Send progress update with current file
+            await progress_callback(files_processed, len(all_functions_data), total_files, current_file)
 
             logger.debug("Processing files", 
                         processed=files_processed, 

backend/tests/test_anonymous_indexing.py

@@ -199,9 +199,9 @@ def test_update_status_completed_with_stats(self, job_manager, mock_redis):
         })
 
         stats = JobStats(
-            files_indexed=100,
-            functions_found=500,
-            time_taken_seconds=45.5
+            files_processed=100,
+            functions_indexed=500,
+            indexing_time_seconds=45.5
         )
 
         result = job_manager.update_status(
@@ -259,13 +259,13 @@ def test_job_progress_none_excluded(self):
     def test_job_stats_to_dict(self):
         """JobStats converts to dict correctly."""
         stats = JobStats(
-            files_indexed=100,
-            functions_found=500,
-            time_taken_seconds=45.5
+            files_processed=100,
+            functions_indexed=500,
+            indexing_time_seconds=45.5
         )
         d = stats.to_dict()
-        assert d["files_indexed"] == 100
-        assert d["time_taken_seconds"] == 45.5
+        assert d["files_processed"] == 100
+        assert d["indexing_time_seconds"] == 45.5
 
 
 # =============================================================================
@@ -608,9 +608,9 @@ def test_completed_job_returns_repo_id(self, mock_job_class, client):
             "created_at": "2024-01-01T00:00:00Z",
             "updated_at": "2024-01-01T00:01:00Z",
             "stats": {
-                "files_indexed": 100,
-                "functions_found": 500,
-                "time_taken_seconds": 45.2
+                "files_processed": 100,
+                "functions_indexed": 500,
+                "indexing_time_seconds": 45.2
             }
         }
         mock_job_class.return_value = mock_job_manager
@@ -621,7 +621,7 @@ def test_completed_job_returns_repo_id(self, mock_job_class, client):
         data = response.json()
         assert data["status"] == "completed"
         assert data["repo_id"] == "anon_idx_test123456"
-        assert data["stats"]["files_indexed"] == 100
+        assert data["stats"]["files_processed"] == 100
 
     @patch('routes.playground.AnonymousIndexingJob')
     def test_failed_job_returns_error(self, mock_job_class, client):

backend/tests/test_ws_playground.py

@@ -66,7 +66,7 @@ def test_websocket_handles_already_completed_job(self):
             "job_id": "idx_test123",
             "status": "completed",
             "repo_id": "anon_test123",
-            "stats": {"files_indexed": 100, "functions_found": 500}
+            "stats": {"files_processed": 100, "functions_indexed": 500}
         })
 
         with patch('routes.ws_playground.redis_client', mock_redis):
@@ -172,9 +172,9 @@ def test_completed_event_includes_stats(self):
         job_manager = AnonymousIndexingJob(mock_redis)
 
         stats = JobStats(
-            files_indexed=100,
-            functions_found=500,
-            time_taken_seconds=45.2
+            files_processed=100,
+            functions_indexed=500,
+            indexing_time_seconds=45.2
         )
 
         job_manager.update_status(
@@ -189,7 +189,7 @@ def test_completed_event_includes_stats(self):
 
         assert event_data["type"] == "completed"
         assert event_data["repo_id"] == "anon_test123"
-        assert event_data["stats"]["functions_found"] == 500
+        assert event_data["stats"]["functions_indexed"] == 500
 
     def test_processing_status_skips_duplicate_publish(self):
         """PROCESSING status should not publish (handled by update_progress)."""

frontend/src/components/playground/HeroPlayground.tsx

@@ -169,11 +169,13 @@ export function HeroPlayground({
   }, [state]);
 
   // Determine visibility states
+  const hasError = state.status === 'error' || wsHasError;
   const showDemoSelector = mode === 'demo';
-  const showUrlInput = mode === 'custom' && !['indexing', 'ready'].includes(state.status) && !showCelebration && !wsIsCompleted;
-  const showValidation = mode === 'custom' && ['validating', 'valid', 'invalid'].includes(state.status) && !showCelebration;
-  const showIndexing = mode === 'custom' && state.status === 'indexing' && !showCelebration && !wsIsCompleted;
+  const showUrlInput = mode === 'custom' && !['indexing', 'ready', 'error'].includes(state.status) && !showCelebration && !wsIsCompleted && !wsHasError;
+  const showValidation = mode === 'custom' && ['validating', 'valid', 'invalid'].includes(state.status) && !showCelebration && !hasError;
+  const showIndexing = mode === 'custom' && state.status === 'indexing' && !showCelebration && !wsIsCompleted && !wsHasError;
   const showReady = (mode === 'custom' && state.status === 'ready') || (wsIsCompleted && !showCelebration);
+  const showError = mode === 'custom' && hasError && !showCelebration;
   const isSearchDisabled = mode === 'custom' && state.status !== 'ready' && !wsIsCompleted;
 
   // Can search?
@@ -350,6 +352,34 @@ export function HeroPlayground({
             </button>
           </motion.div>
         )}
+
+        {/* Error State - Inside AnimatePresence to prevent simultaneous renders */}
+        {showError && (
+          <motion.div
+            key="error"
+            initial={{ opacity: 0, y: 10 }}
+            animate={{ opacity: 1, y: 0 }}
+            exit={{ opacity: 0, y: -10 }}
+            className="mb-4 px-4 py-3 rounded-xl bg-amber-500/10 border border-amber-500/20"
+          >
+            <div className="flex items-start gap-3">
+              <span className="text-amber-400 text-lg">⚠️</span>
+              <div className="flex-1">
+                <p className="text-amber-200 text-sm font-medium mb-1">Something went wrong</p>
+                <p className="text-zinc-400 text-sm">
+                  {state.status === 'error' ? state.message : wsError || 'An unexpected error occurred'}
+                </p>
+                <button 
+                  type="button"
+                  onClick={handleIndexAnother}
+                  className="mt-3 text-sm text-amber-400 hover:text-amber-300 transition-colors font-medium"
+                >
+                  ← Try again with a different repository
+                </button>
+              </div>
+            </div>
+          </motion.div>
+        )}
       </AnimatePresence>
 
       {/* Search Box */}
@@ -407,26 +437,6 @@ export function HeroPlayground({
         </>
       )}
 
-      {/* Error State */}
-      {(state.status === 'error' || wsHasError) && (
-        <motion.div
-          initial={{ opacity: 0, y: 10 }}
-          animate={{ opacity: 1, y: 0 }}
-          className="mt-4 px-4 py-3 rounded-xl bg-red-500/10 border border-red-500/20"
-        >
-          <p className="text-red-300 text-sm">
-            {state.status === 'error' ? state.message : wsError}
-          </p>
-          <button 
-            type="button"
-            onClick={handleIndexAnother}
-            className="mt-2 text-xs text-red-400 hover:text-red-300 transition-colors"
-          >
-            Try again
-          </button>
-        </motion.div>
-      )}
-
       {/* Upgrade CTA (when limit reached) */}
       {remaining <= 0 && (
         <div className="mt-6 p-4 rounded-xl bg-gradient-to-r from-indigo-600/20 to-purple-600/20 border border-indigo-500/30 text-center">

frontend/src/components/playground/IndexingProgress.tsx

@@ -267,10 +267,19 @@ export function IndexingProgress({
 }: IndexingProgressProps) {
   const { percent, filesProcessed, filesTotal, currentFile, functionsFound } = progress;
 
-  // Estimate remaining time
+  // Safe values to prevent NaN/division by zero
+  const safePercent = Number.isFinite(percent) ? Math.max(0, Math.min(100, percent)) : 0;
+  const safeFilesProcessed = Number.isFinite(filesProcessed) ? Math.max(0, filesProcessed) : 0;
+  const safeFilesTotal = Number.isFinite(filesTotal) ? Math.max(0, filesTotal) : 0;
+  const safeFunctionsFound = Number.isFinite(functionsFound) ? Math.max(0, functionsFound) : 0;
+
+  // Is this the initial "starting" state?
+  const isStarting = safeFilesTotal === 0 || (safePercent === 0 && safeFilesProcessed === 0);
+
+  // Estimate remaining time (only when we have real data)
   const estimatedRemaining = (() => {
-    if (percent <= 0 || filesProcessed <= 0) return null;
-    const remainingFiles = Math.ceil((filesProcessed / percent) * (100 - percent));
+    if (isStarting || safePercent <= 0 || safeFilesProcessed <= 0) return null;
+    const remainingFiles = Math.ceil((safeFilesProcessed / safePercent) * (100 - safePercent));
     return Math.max(1, Math.ceil(remainingFiles * 0.15));
   })();
 
@@ -287,7 +296,7 @@ export function IndexingProgress({
       initial={{ opacity: 0, y: 20 }}
       animate={{ opacity: 1, y: 0 }}
       role="status"
-      aria-label={`Indexing ${repoName || 'repository'}: ${percent}% complete`}
+      aria-label={`Indexing ${repoName || 'repository'}: ${safePercent}% complete`}
     >
       {/* Header */}
       <div className="px-5 py-4 border-b border-zinc-800/50 bg-gradient-to-r from-zinc-900 to-zinc-900/50">
@@ -308,12 +317,12 @@ export function IndexingProgress({
           </div>
           <motion.span 
             className="text-3xl font-bold text-indigo-400"
-            key={percent}
+            key={safePercent}
             initial={{ scale: 1.3 }}
             animate={{ scale: 1 }}
             transition={{ type: 'spring', stiffness: 300 }}
           >
-            {percent}%
+            {safePercent}%
           </motion.span>
         </div>
 
@@ -323,16 +332,29 @@ export function IndexingProgress({
 
       {/* Progress bar */}
       <div className="px-5 py-4">
-        <GlowingProgress value={percent} />
+        <GlowingProgress value={safePercent} />
       </div>
 
       {/* Stats grid */}
       <div className="px-5 py-4 bg-zinc-900/50 border-t border-b border-zinc-800/50">
         <div className="grid grid-cols-4 gap-4">
-          <StatCard label="Files" value={`${filesProcessed}/${filesTotal}`} />
-          <StatCard label="Functions" value={functionsFound} highlight />
-          <StatCard label="Remaining" value={estimatedRemaining ? `~${estimatedRemaining}s` : '—'} />
-          <StatCard label="Speed" value={filesProcessed > 0 ? `${Math.round(functionsFound / filesProcessed)}/file` : '—'} />
+          <StatCard 
+            label="Files" 
+            value={isStarting ? 'Starting...' : `${safeFilesProcessed}/${safeFilesTotal}`} 
+          />
+          <StatCard 
+            label="Functions" 
+            value={isStarting ? '—' : safeFunctionsFound} 
+            highlight={!isStarting && safeFunctionsFound > 0} 
+          />
+          <StatCard 
+            label="Remaining" 
+            value={estimatedRemaining ? `~${estimatedRemaining}s` : '—'} 
+          />
+          <StatCard 
+            label="Speed" 
+            value={safeFilesProcessed > 0 ? `${Math.round(safeFunctionsFound / safeFilesProcessed)}/file` : '—'} 
+          />
         </div>
       </div>
 

frontend/src/hooks/useIndexingWebSocket.ts

@@ -248,8 +248,18 @@ export function useIndexingWebSocket(
     if (jobId) {
       connect(jobId);
     } else {
+      // Only cleanup connection, DON'T reset state!
+      // This preserves completedStats when jobId becomes null after completion
       cleanup();
-      setState(INITIAL_STATE);
+      // Only reset if we were never completed (e.g., user navigated away during indexing)
+      setState(prev => {
+        if (prev.phase === 'completed') {
+          // Keep completed state - just disconnect
+          return { ...prev, connectionState: 'disconnected' };
+        }
+        // Reset if we were mid-indexing (user cancelled, navigated away, etc.)
+        return INITIAL_STATE;
+      });
     }
     return cleanup;
   }, [jobId, connect, cleanup]);