fix(ws): critical bugs in WebSocket endpoint

- Fix: Must accept() before close() - was causing runtime errors
- Fix: Check job exists before subscribing (return 4404)
- Fix: Handle race condition - job may complete before WS connects
- Fix: Remove unused INITIAL_TIMEOUT_SECONDS constant
- Add: Idle timeout (120s) to prevent zombie connections
- Add: 6 new tests for edge cases (13 total)

Review by CTO caught these before production 🙏

Author: DevanshuNEU

backend/routes/ws_playground.py

@@ -5,7 +5,7 @@
 a smooth streaming experience instead of polling every 2 seconds.
 
 Channel format: job:{job_id}:events
-Message types: started, progress, completed, error
+Message types: connected, cloning, progress, completed, error
 """
 import json
 import asyncio
@@ -15,13 +15,14 @@
 
 from dependencies import redis_client
 from services.observability import logger
+from services.anonymous_indexer import AnonymousIndexingJob
 
 
-# How long to wait for first message before giving up
-INITIAL_TIMEOUT_SECONDS = 30
+# How long between messages before sending a ping
+PING_INTERVAL_SECONDS = 30
 
-# How long between messages before assuming job is dead
-MESSAGE_TIMEOUT_SECONDS = 60
+# How long to wait for any activity before closing
+IDLE_TIMEOUT_SECONDS = 120
 
 
 async def websocket_playground_index(websocket: WebSocket, job_id: str):
@@ -35,59 +36,109 @@ async def websocket_playground_index(websocket: WebSocket, job_id: str):
     No auth required - job_id is an unguessable UUID that acts as
     a bearer token. Only the session that created the job knows it.
     """
+    # Validate job_id format (basic sanity check)
+    if not job_id or len(job_id) < 10:
+        # Must accept before we can close with a reason
+        await websocket.accept()
+        await websocket.close(code=4400, reason="Invalid job ID")
+        return
+    
     # Validate we have Redis (required for pub/sub)
     if not redis_client:
         logger.error("WebSocket failed - no Redis connection")
+        await websocket.accept()
         await websocket.close(code=4500, reason="Service unavailable")
         return
 
-    # Validate job_id format (basic sanity check)
-    if not job_id or len(job_id) < 10:
-        await websocket.close(code=4400, reason="Invalid job ID")
-        return
+    # Check if job exists before subscribing
+    job_manager = AnonymousIndexingJob(redis_client)
+    job = job_manager.get_job(job_id)
 
-    channel = f"job:{job_id}:events"
+    if not job:
+        await websocket.accept()
+        await websocket.close(code=4404, reason="Job not found")
+        return
 
     # Accept the WebSocket connection
     await websocket.accept()
-    logger.info("WebSocket connected", job_id=job_id[:12], channel=channel)
+    logger.info("WebSocket connected", job_id=job_id[:12])
 
-    # Set up Redis pub/sub
+    # Handle race condition: job might already be complete
+    job_status = job.get("status")
+    if job_status == "completed":
+        await websocket.send_json({
+            "type": "completed",
+            "job_id": job_id,
+            "repo_id": job.get("repo_id"),
+            "stats": job.get("stats"),
+            "message": "Indexing already complete"
+        })
+        await websocket.close()
+        return
+    elif job_status == "failed":
+        await websocket.send_json({
+            "type": "error",
+            "job_id": job_id,
+            "error": job.get("error"),
+            "message": job.get("error_message", "Indexing failed"),
+            "recoverable": False
+        })
+        await websocket.close()
+        return
+    
+    channel = f"job:{job_id}:events"
     pubsub = redis_client.pubsub()
 
     try:
         # Subscribe to job's event channel
         await asyncio.to_thread(pubsub.subscribe, channel)
         logger.debug("Subscribed to channel", channel=channel)
 
-        # Send initial ack so client knows we're connected
+        # Send initial ack with current state
         await websocket.send_json({
             "type": "connected",
             "job_id": job_id,
+            "current_status": job_status,
             "message": "Listening for indexing events"
         })
 
         # Listen for messages
+        last_activity = asyncio.get_event_loop().time()
+        
         while True:
-            # Check for new message (non-blocking with timeout)
+            current_time = asyncio.get_event_loop().time()
+            
+            # Check for idle timeout
+            if current_time - last_activity > IDLE_TIMEOUT_SECONDS:
+                logger.warning("WebSocket idle timeout", job_id=job_id[:12])
+                await websocket.send_json({
+                    "type": "error",
+                    "message": "Connection timed out - no activity"
+                })
+                break
+            
+            # Check for new message (non-blocking with short timeout)
             message = await asyncio.to_thread(
                 pubsub.get_message,
                 ignore_subscribe_messages=True,
-                timeout=MESSAGE_TIMEOUT_SECONDS
+                timeout=PING_INTERVAL_SECONDS
             )
 
             if message is None:
-                # Timeout - check if client is still connected
+                # No message - send ping to keep connection alive
                 try:
                     await websocket.send_json({"type": "ping"})
                 except Exception:
-                    logger.debug("Client disconnected during timeout", job_id=job_id[:12])
+                    logger.debug("Client disconnected during ping", job_id=job_id[:12])
                     break
                 continue
 
             if message["type"] != "message":
                 continue
 
+            # Got a message - reset activity timer
+            last_activity = current_time
+            
             # Parse and forward the event
             try:
                 event_data = json.loads(message["data"])

backend/tests/test_ws_playground.py

@@ -10,9 +10,6 @@
 from fastapi.testclient import TestClient
 from fastapi.websockets import WebSocket
 
-# We'll test the WebSocket handler directly since TestClient
-# doesn't support async WebSocket testing well
-
 
 class TestWebSocketPlayground:
     """Test suite for playground WebSocket endpoint."""
@@ -23,27 +20,87 @@ def test_websocket_rejects_invalid_job_id(self):
 
         client = TestClient(app)
 
-        # Job ID too short
-        with pytest.raises(Exception):
-            with client.websocket_connect("/api/v1/ws/playground/abc"):
-                pass
+        # Job ID too short - should accept then close with 4400
+        with client.websocket_connect("/api/v1/ws/playground/abc") as ws:
+            # Connection will be closed immediately
+            pass
 
     def test_websocket_rejects_when_no_redis(self):
         """Should close gracefully if Redis is unavailable."""
         from routes.ws_playground import websocket_playground_index
 
-        # Mock WebSocket
         mock_ws = AsyncMock(spec=WebSocket)
 
-        # Patch redis_client to None
         with patch('routes.ws_playground.redis_client', None):
             asyncio.run(websocket_playground_index(mock_ws, "idx_abc123def456"))
 
-        # Should close with service unavailable
+        # Should accept first, then close
+        mock_ws.accept.assert_called_once()
         mock_ws.close.assert_called_once()
         call_args = mock_ws.close.call_args
         assert call_args[1]['code'] == 4500
 
+    def test_websocket_rejects_nonexistent_job(self):
+        """Should close with 4404 if job doesn't exist."""
+        from routes.ws_playground import websocket_playground_index
+        
+        mock_ws = AsyncMock(spec=WebSocket)
+        mock_redis = MagicMock()
+        mock_redis.get.return_value = None  # Job not found
+        
+        with patch('routes.ws_playground.redis_client', mock_redis):
+            asyncio.run(websocket_playground_index(mock_ws, "idx_nonexistent"))
+        
+        mock_ws.accept.assert_called_once()
+        mock_ws.close.assert_called_once()
+        call_args = mock_ws.close.call_args
+        assert call_args[1]['code'] == 4404
+
+    def test_websocket_handles_already_completed_job(self):
+        """If job is already complete, send completion and close."""
+        from routes.ws_playground import websocket_playground_index
+        
+        mock_ws = AsyncMock(spec=WebSocket)
+        mock_redis = MagicMock()
+        mock_redis.get.return_value = json.dumps({
+            "job_id": "idx_test123",
+            "status": "completed",
+            "repo_id": "anon_test123",
+            "stats": {"files_indexed": 100, "functions_found": 500}
+        })
+        
+        with patch('routes.ws_playground.redis_client', mock_redis):
+            asyncio.run(websocket_playground_index(mock_ws, "idx_test123"))
+        
+        # Should send completed event immediately
+        mock_ws.send_json.assert_called_once()
+        sent_data = mock_ws.send_json.call_args[0][0]
+        assert sent_data["type"] == "completed"
+        assert sent_data["repo_id"] == "anon_test123"
+        mock_ws.close.assert_called_once()
+
+    def test_websocket_handles_already_failed_job(self):
+        """If job already failed, send error and close."""
+        from routes.ws_playground import websocket_playground_index
+        
+        mock_ws = AsyncMock(spec=WebSocket)
+        mock_redis = MagicMock()
+        mock_redis.get.return_value = json.dumps({
+            "job_id": "idx_test123",
+            "status": "failed",
+            "error": "clone_failed",
+            "error_message": "Repository not found"
+        })
+        
+        with patch('routes.ws_playground.redis_client', mock_redis):
+            asyncio.run(websocket_playground_index(mock_ws, "idx_test123"))
+        
+        mock_ws.send_json.assert_called_once()
+        sent_data = mock_ws.send_json.call_args[0][0]
+        assert sent_data["type"] == "error"
+        assert sent_data["error"] == "clone_failed"
+        mock_ws.close.assert_called_once()
+
 
 class TestPubSubIntegration:
     """Test Redis Pub/Sub event publishing."""
@@ -52,7 +109,6 @@ def test_publish_event_called_on_status_update(self):
         """Verify events are published when status changes."""
         from services.anonymous_indexer import AnonymousIndexingJob, JobStatus
 
-        # Create mock Redis
         mock_redis = MagicMock()
         mock_redis.get.return_value = json.dumps({
             "job_id": "idx_test123",
@@ -61,11 +117,8 @@ def test_publish_event_called_on_status_update(self):
         })
 
         job_manager = AnonymousIndexingJob(mock_redis)
-        
-        # Update status to cloning
         job_manager.update_status("idx_test123", JobStatus.CLONING)
 
-        # Verify publish was called
         mock_redis.publish.assert_called()
         call_args = mock_redis.publish.call_args
 
@@ -88,7 +141,6 @@ def test_progress_event_published_with_file_info(self):
 
         job_manager = AnonymousIndexingJob(mock_redis)
 
-        # Update progress
         job_manager.update_progress(
             job_id="idx_test123",
             files_processed=25,
@@ -132,7 +184,6 @@ def test_completed_event_includes_stats(self):
             repo_id="anon_test123"
         )
 
-        # Verify publish was called with completed event
         call_args = mock_redis.publish.call_args
         event_data = json.loads(call_args[0][1])
 
@@ -152,13 +203,38 @@ def test_processing_status_skips_duplicate_publish(self):
 
         job_manager = AnonymousIndexingJob(mock_redis)
 
-        # Call update_status with PROCESSING directly
         progress = JobProgress(files_total=100, files_processed=0, functions_found=0)
         job_manager.update_status("idx_test123", JobStatus.PROCESSING, progress=progress)
 
         # Should NOT have published (returns early for PROCESSING)
         mock_redis.publish.assert_not_called()
 
+    def test_failed_event_includes_recoverable_flag(self):
+        """Failed events should indicate if error is recoverable."""
+        from services.anonymous_indexer import AnonymousIndexingJob, JobStatus
+        
+        mock_redis = MagicMock()
+        mock_redis.get.return_value = json.dumps({
+            "job_id": "idx_test123",
+            "status": "processing"
+        })
+        
+        job_manager = AnonymousIndexingJob(mock_redis)
+        
+        # Clone failures are not recoverable
+        job_manager.update_status(
+            "idx_test123",
+            JobStatus.FAILED,
+            error="clone_failed",
+            error_message="Repo not found"
+        )
+        
+        call_args = mock_redis.publish.call_args
+        event_data = json.loads(call_args[0][1])
+        
+        assert event_data["type"] == "failed"
+        assert event_data["recoverable"] == False
+
 
 class TestEventFormats:
     """Verify event message formats match frontend expectations."""
@@ -179,6 +255,52 @@ def test_event_types_are_strings(self):
         call_args = mock_redis.publish.call_args
         event_data = json.loads(call_args[0][1])
 
-        # Type should be string "cloning", not JobStatus.CLONING
         assert isinstance(event_data["type"], str)
         assert event_data["type"] == "cloning"
+
+    def test_progress_percent_is_integer(self):
+        """Progress percent should be an integer 0-100."""
+        from services.anonymous_indexer import AnonymousIndexingJob
+        
+        mock_redis = MagicMock()
+        mock_redis.get.return_value = json.dumps({
+            "job_id": "idx_test123",
+            "status": "processing"
+        })
+        
+        job_manager = AnonymousIndexingJob(mock_redis)
+        
+        job_manager.update_progress(
+            job_id="idx_test123",
+            files_processed=33,
+            functions_found=100,
+            files_total=100,
+            current_file="test.py"
+        )
+        
+        progress_call = mock_redis.publish.call_args_list[0]
+        event_data = json.loads(progress_call[0][1])
+        
+        assert isinstance(event_data["percent"], int)
+        assert event_data["percent"] == 33
+
+    def test_all_events_include_job_id(self):
+        """All events should include job_id for client correlation."""
+        from services.anonymous_indexer import AnonymousIndexingJob, JobStatus
+        
+        mock_redis = MagicMock()
+        mock_redis.get.return_value = json.dumps({
+            "job_id": "idx_test123",
+            "status": "queued"
+        })
+        
+        job_manager = AnonymousIndexingJob(mock_redis)
+        
+        # Test different event types
+        job_manager.update_status("idx_test123", JobStatus.CLONING)
+        
+        call_args = mock_redis.publish.call_args
+        event_data = json.loads(call_args[0][1])
+        
+        assert "job_id" in event_data
+        assert event_data["job_id"] == "idx_test123"