feat: add WebSocket support for real-time scraper progress updates

backend/app/api/routes/episode.py

@@ -136,15 +136,36 @@ async def step_episode(request: StepRequest) -> StepResponse:
     logger.info(f"Step in episode {request.episode_id}: {request.action.action_type}")
 
     env = get_environment(request.episode_id)
+    
+    # Get WebSocket manager for real-time updates
+    from app.api.routes.websocket import get_connection_manager
+    ws_manager = get_connection_manager()
 
     try:
         observation, reward, reward_breakdown, terminated, truncated, info = await env.step(
             request.action
         )
+        
+        # Send real-time progress update via WebSocket
+        await ws_manager.send_progress_update(
+            episode_id=request.episode_id,
+            step=observation.step_number,
+            action_type=request.action.action_type.value,
+            reward=reward,
+            progress=observation.extraction_progress,
+            message=f"Executed {request.action.action_type.value}",
+        )
 
         # Clean up if episode is done
         if terminated or truncated:
             logger.info(f"Episode {request.episode_id} completed")
+            state = env.get_state()
+            await ws_manager.send_completion(
+                episode_id=request.episode_id,
+                success=terminated and not truncated,
+                total_reward=state.get("total_reward", 0.0),
+                extracted_data=state.get("extracted_data", {}),
+            )
 
         return StepResponse(
             observation=observation,
@@ -156,6 +177,11 @@ async def step_episode(request: StepRequest) -> StepResponse:
         )
     except Exception as e:
         logger.error(f"Step failed: {e}")
+        await ws_manager.send_error(
+            episode_id=request.episode_id,
+            error=str(e),
+            details={"action_type": request.action.action_type.value},
+        )
         raise HTTPException(
             status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
             detail=f"Step execution failed: {str(e)}",

backend/app/api/routes/websocket.py

@@ -0,0 +1,198 @@
+"""WebSocket support for real-time scraper updates."""
+
+import asyncio
+import json
+import logging
+from typing import Any
+
+from fastapi import APIRouter, WebSocket, WebSocketDisconnect
+from fastapi.websockets import WebSocketState
+
+logger = logging.getLogger(__name__)
+
+router = APIRouter(prefix="/ws", tags=["WebSocket"])
+
+# Store active WebSocket connections by episode_id
+_active_connections: dict[str, list[WebSocket]] = {}
+
+
+class ConnectionManager:
+    """Manage WebSocket connections for real-time updates."""
+
+    def __init__(self):
+        self.active_connections: dict[str, list[WebSocket]] = {}
+
+    async def connect(self, websocket: WebSocket, episode_id: str):
+        """Connect a new WebSocket client."""
+        await websocket.accept()
+        if episode_id not in self.active_connections:
+            self.active_connections[episode_id] = []
+        self.active_connections[episode_id].append(websocket)
+        logger.info(f"WebSocket connected for episode {episode_id}")
+
+    def disconnect(self, websocket: WebSocket, episode_id: str):
+        """Disconnect a WebSocket client."""
+        if episode_id in self.active_connections:
+            if websocket in self.active_connections[episode_id]:
+                self.active_connections[episode_id].remove(websocket)
+            if not self.active_connections[episode_id]:
+                del self.active_connections[episode_id]
+        logger.info(f"WebSocket disconnected for episode {episode_id}")
+
+    async def send_personal_message(self, message: dict[str, Any], websocket: WebSocket):
+        """Send a message to a specific client."""
+        try:
+            if websocket.client_state == WebSocketState.CONNECTED:
+                await websocket.send_json(message)
+        except Exception as e:
+            logger.error(f"Error sending personal message: {e}")
+
+    async def broadcast(self, message: dict[str, Any], episode_id: str):
+        """Broadcast a message to all clients watching an episode."""
+        if episode_id not in self.active_connections:
+            return
+
+        disconnected = []
+        for connection in self.active_connections[episode_id]:
+            try:
+                if connection.client_state == WebSocketState.CONNECTED:
+                    await connection.send_json(message)
+                else:
+                    disconnected.append(connection)
+            except Exception as e:
+                logger.error(f"Error broadcasting to client: {e}")
+                disconnected.append(connection)
+
+        # Clean up disconnected clients
+        for conn in disconnected:
+            self.disconnect(conn, episode_id)
+
+    async def send_progress_update(
+        self,
+        episode_id: str,
+        step: int,
+        action_type: str,
+        reward: float,
+        progress: float,
+        message: str | None = None,
+    ):
+        """Send a progress update for an episode."""
+        update = {
+            "type": "progress",
+            "episode_id": episode_id,
+            "step": step,
+            "action_type": action_type,
+            "reward": reward,
+            "progress": progress,
+            "message": message,
+            "timestamp": asyncio.get_event_loop().time(),
+        }
+        await self.broadcast(update, episode_id)
+
+    async def send_error(self, episode_id: str, error: str, details: dict[str, Any] | None = None):
+        """Send an error message."""
+        message = {
+            "type": "error",
+            "episode_id": episode_id,
+            "error": error,
+            "details": details or {},
+            "timestamp": asyncio.get_event_loop().time(),
+        }
+        await self.broadcast(message, episode_id)
+
+    async def send_completion(
+        self,
+        episode_id: str,
+        success: bool,
+        total_reward: float,
+        extracted_data: dict[str, Any],
+    ):
+        """Send a completion notification."""
+        message = {
+            "type": "completion",
+            "episode_id": episode_id,
+            "success": success,
+            "total_reward": total_reward,
+            "extracted_data": extracted_data,
+            "timestamp": asyncio.get_event_loop().time(),
+        }
+        await self.broadcast(message, episode_id)
+
+
+# Global connection manager
+manager = ConnectionManager()
+
+
+@router.websocket("/episode/{episode_id}")
+async def websocket_episode(websocket: WebSocket, episode_id: str):
+    """
+    WebSocket endpoint for receiving real-time updates about an episode.
+    
+    Clients can connect to this endpoint to receive updates about:
+    - Action execution progress
+    - Reward changes
+    - Extraction progress
+    - Errors
+    - Episode completion
+    
+    Args:
+        websocket: WebSocket connection
+        episode_id: ID of the episode to watch
+    """
+    await manager.connect(websocket, episode_id)
+    
+    try:
+        # Send initial connection confirmation
+        await manager.send_personal_message(
+            {
+                "type": "connected",
+                "episode_id": episode_id,
+                "message": f"Connected to episode {episode_id}",
+            },
+            websocket,
+        )
+        
+        # Keep connection alive and handle incoming messages
+        while True:
+            try:
+                # Receive messages from client (e.g., subscription updates)
+                data = await asyncio.wait_for(
+                    websocket.receive_text(),
+                    timeout=30.0,  # 30 second timeout
+                )
+                
+                try:
+                    message = json.loads(data)
+                    
+                    # Handle ping/pong for keep-alive
+                    if message.get("type") == "ping":
+                        await manager.send_personal_message(
+                            {"type": "pong", "timestamp": asyncio.get_event_loop().time()},
+                            websocket,
+                        )
+                    
+                except json.JSONDecodeError:
+                    logger.warning(f"Invalid JSON received: {data}")
+                    
+            except asyncio.TimeoutError:
+                # Send a ping to check if client is still connected
+                try:
+                    await manager.send_personal_message(
+                        {"type": "ping", "timestamp": asyncio.get_event_loop().time()},
+                        websocket,
+                    )
+                except Exception:
+                    # Client disconnected
+                    break
+                    
+    except WebSocketDisconnect:
+        logger.info(f"Client disconnected from episode {episode_id}")
+    except Exception as e:
+        logger.error(f"WebSocket error for episode {episode_id}: {e}")
+    finally:
+        manager.disconnect(websocket, episode_id)
+
+
+def get_connection_manager() -> ConnectionManager:
+    """Get the global connection manager instance."""
+    return manager

backend/app/main.py

@@ -137,6 +137,10 @@ def create_app() -> FastAPI:
     # Import and include providers router
     from app.api.routes import providers
     app.include_router(providers.router, prefix=api_prefix, tags=["Providers"])
+    
+    # Import and include WebSocket router
+    from app.api.routes import websocket
+    app.include_router(websocket.router, tags=["WebSocket"])
 
     # Serve static files (frontend build)
     static_dir = Path(__file__).parent.parent / "static"