import asyncio import sys import time from typing import Dict, List, Optional, Set, Any from fastapi import WebSocket from crawler.config import CrawlerConfig from crawler.engine import Crawler class CrawlSessionManager: """Manages active crawl sessions, WebSocket broadcasting, metrics, and error aggregation.""" def __init__(self): self.active_connections: Set[WebSocket] = set() self.crawler: Optional[Crawler] = None self.task: Optional[asyncio.Task] = None self.status: str = "idle" # "idle" | "running" | "stopping" | "finished" | "error" self.status_message: str = "Ready to start" self.start_time: Optional[float] = None self.end_time: Optional[float] = None self.current_iteration: int = 1 self.total_iterations: int = 1 self.iteration_mode: str = "all" self.max_visits_per_iteration: int = 10 self.stats: Dict[str, int] = { "internal": 0, "external": 0, "cookies_accepted": 0, } self.queue_size: int = 0 self.visited_count: int = 0 self.total_visited_all: int = 0 self.total_errors_all: int = 0 self.tor_identity: Optional[Dict[str, Any]] = None self.errors: List[Dict[str, Any]] = [] self.logs: List[Dict[str, Any]] = [] self._log_counter: int = 0 self._error_counter: int = 0 async def connect_ws(self, websocket: WebSocket): """Registers a new WebSocket connection and transmits initial state snapshot.""" await websocket.accept() self.active_connections.add(websocket) # Send current state snapshot immediately upon connection snapshot = self.get_snapshot() await websocket.send_json({ "type": "init", "data": snapshot, }) def disconnect_ws(self, websocket: WebSocket): """Removes a disconnected WebSocket.""" self.active_connections.discard(websocket) async def broadcast(self, message: Dict[str, Any]): """Dispatches an event payload to all connected WebSockets.""" disconnected = set() for ws in self.active_connections: try: await ws.send_json(message) except Exception: disconnected.add(ws) for ws in disconnected: self.active_connections.discard(ws) def get_snapshot(self) -> Dict[str, Any]: """Generates a complete snapshot of current crawler session.""" elapsed = 0.0 if self.start_time: if self.end_time: elapsed = self.end_time - self.start_time else: elapsed = time.time() - self.start_time return { "status": self.status, "status_message": self.status_message, "elapsed_seconds": round(elapsed, 1), "current_iteration": self.current_iteration, "total_iterations": self.total_iterations, "iteration_mode": self.iteration_mode, "max_visits_per_iteration": self.max_visits_per_iteration, "stats": self.stats, "queue_size": self.queue_size, "visited_count": self.visited_count, "total_visited_all": self.total_visited_all, "total_errors_all": max(self.total_errors_all, len(self.errors)), "tor_identity": self.tor_identity, "error_count": len(self.errors), "errors": self.errors[-100:], # last 100 errors for quick load "logs": self.logs[-500:], # last 500 logs for quick load } async def handle_crawler_event(self, event: Dict[str, Any]): """Processes events dispatched by the core crawler engine.""" event_type = event.get("type") if event_type == "log": self._log_counter += 1 log_entry = { "id": self._log_counter, "message": event.get("message", ""), "level": event.get("level", "info"), "timestamp": event.get("timestamp", time.time()), } self.logs.append(log_entry) if len(self.logs) > 2000: self.logs = self.logs[-1500:] await self.broadcast({"type": "log", "data": log_entry}) elif event_type == "stats": self.stats = event.get("stats", self.stats) self.queue_size = event.get("queue_size", self.queue_size) self.visited_count = event.get("visited_count", self.visited_count) if "total_visited_all" in event: self.total_visited_all = event["total_visited_all"] else: self.total_visited_all = max(self.total_visited_all, self.visited_count) if "total_errors_all" in event: self.total_errors_all = event["total_errors_all"] else: self.total_errors_all = len(self.errors) await self.broadcast({ "type": "stats", "data": { "stats": self.stats, "queue_size": self.queue_size, "visited_count": self.visited_count, "error_count": event.get("error_count", 0), "total_visited_all": self.total_visited_all, "total_errors_all": self.total_errors_all, }, }) elif event_type == "tor_identity": self.tor_identity = event.get("data", {}) await self.broadcast({ "type": "tor_identity", "data": self.tor_identity, }) elif event_type == "error": self._error_counter += 1 error_entry = { "id": self._error_counter, "url": event.get("url", ""), "source_url": event.get("source_url") or "Root / Direct Input", "error": event.get("error", "Unknown error"), "timestamp": event.get("timestamp", time.time()), } self.errors.append(error_entry) self.total_errors_all = len(self.errors) await self.broadcast({ "type": "error", "data": { **error_entry, "total_errors_all": self.total_errors_all, }, }) elif event_type == "iteration": self.current_iteration = event.get("iteration", 1) self.total_iterations = event.get("total", 1) if "iteration_mode" in event: self.iteration_mode = event["iteration_mode"] if "max_visits_per_iteration" in event: self.max_visits_per_iteration = event["max_visits_per_iteration"] if "total_visited_all" in event: self.total_visited_all = event["total_visited_all"] if "total_errors_all" in event: self.total_errors_all = event["total_errors_all"] await self.broadcast({ "type": "iteration", "data": { "current": self.current_iteration, "total": self.total_iterations, "iteration_mode": self.iteration_mode, "max_visits_per_iteration": self.max_visits_per_iteration, "total_visited_all": self.total_visited_all, "total_errors_all": self.total_errors_all, }, }) elif event_type == "status": self.status = event.get("status", self.status) self.status_message = event.get("message", "") if "total_visited_all" in event: self.total_visited_all = event["total_visited_all"] if "total_errors_all" in event: self.total_errors_all = event["total_errors_all"] if self.status in ["finished", "stopped", "error"]: self.end_time = time.time() await self.broadcast({ "type": "status", "data": { "status": self.status, "message": self.status_message, "elapsed": round((self.end_time or time.time()) - (self.start_time or time.time()), 1), "total_visited_all": self.total_visited_all, "total_errors_all": self.total_errors_all, }, }) async def start_crawl(self, config_data: Dict[str, Any]) -> Dict[str, Any]: """Configures and initiates a new crawl in a background task.""" if self.status == "running": return {"success": False, "message": "Crawler is already running."} # Safeguard: Verify Windows event loop supports subprocesses for Playwright if sys.platform == "win32": loop = asyncio.get_running_loop() if not isinstance(loop, getattr(asyncio, "ProactorEventLoop", ())): err_msg = ( "Current asyncio event loop does not support subprocesses on Windows " "(ProactorEventLoop is required for Playwright). " "Please launch Websurfer using 'python web.py' or ensure Uvicorn uses '--loop asyncio:ProactorEventLoop'." ) return {"success": False, "message": err_msg} # Clear session data self.stats = {"internal": 0, "external": 0, "cookies_accepted": 0} self.queue_size = 0 self.visited_count = 0 self.total_visited_all = 0 self.total_errors_all = 0 self.tor_identity = None self.errors = [] self.logs = [] self._log_counter = 0 self._error_counter = 0 self.start_time = time.time() self.end_time = None self.status = "running" self.status_message = "Initializing crawler..." self.current_iteration = 1 self.total_iterations = int(config_data.get("num_iterations", 1)) # Build configuration object config = CrawlerConfig.from_dict(config_data) self.iteration_mode = config.iteration_mode self.max_visits_per_iteration = config.max_visits_per_iteration # Instantiate crawler with event callback self.crawler = Crawler(config, event_callback=self.handle_crawler_event) async def _run_task(): try: await self.crawler.run() except Exception as e: await self.handle_crawler_event({ "type": "status", "status": "error", "message": f"Execution error: {e}", }) finally: if self.status == "running": self.status = "finished" self.end_time = time.time() await self.broadcast({ "type": "status", "data": { "status": self.status, "message": "Completed", "elapsed": round(self.end_time - self.start_time, 1), }, }) self.task = asyncio.create_task(_run_task()) await self.broadcast({"type": "init", "data": self.get_snapshot()}) return {"success": True, "message": "Crawl started"} async def stop_crawl(self) -> Dict[str, Any]: """Stops the active crawler job.""" if self.status != "running" or not self.crawler: return {"success": False, "message": "No active crawl to stop."} self.status = "stopping" self.status_message = "Cancelling workers..." await self.broadcast({ "type": "status", "data": {"status": "stopping", "message": "Cancelling workers..."}, }) self.crawler.stop() return {"success": True, "message": "Stop signal sent."} def clear_logs(self): """Clears in-memory logs.""" self.logs = [] self._log_counter = 0 # Global singleton manager instance session_manager = CrawlSessionManager()